argument_parser.hpp 9.05 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
22
23
24

namespace migraphx {
namespace driver {
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
25

Paul's avatar
Paul committed
26
27
28
29
30
31
#ifdef MIGRAPHX_USE_CLANG_TIDY
#define MIGRAPHX_DRIVER_STATIC
#else
#define MIGRAPHX_DRIVER_STATIC static
#endif

Paul's avatar
Paul committed
32
template <class T>
Paul's avatar
Paul committed
33
34
35
36
37
using bare = std::remove_cv_t<std::remove_reference_t<T>>;

namespace detail {

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

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

Paul's avatar
Paul committed
54
55
56
template <class T>
struct value_parser
{
Paul's avatar
Paul committed
57
    template <MIGRAPHX_REQUIRES(not std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
58
59
60
61
62
63
64
65
66
67
68
    static T apply(const std::string& x)
    {
        T result;
        std::stringstream ss;
        ss.str(x);
        ss >> result;
        if(ss.fail())
            throw std::runtime_error("Failed to parse: " + x);
        return result;
    }

Paul's avatar
Paul committed
69
    template <MIGRAPHX_REQUIRES(std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
70
71
72
73
74
75
76
77
78
79
    static T apply(const std::string& x)
    {
        std::ptrdiff_t i;
        std::stringstream ss;
        ss.str(x);
        ss >> i;
        if(ss.fail())
            throw std::runtime_error("Failed to parse: " + x);
        return static_cast<T>(i);
    }
Paul's avatar
Paul committed
80
81
82
83
84
85
86
87
88

    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
89
90
91
92
93
94
95
};

struct argument_parser
{
    struct argument
    {
        std::vector<std::string> flags;
Paul's avatar
Paul committed
96
        std::function<bool(argument_parser&, const std::vector<std::string>&)> action{};
Paul's avatar
Paul committed
97
98
99
        std::string type          = "";
        std::string help          = "";
        std::string metavar       = "";
Paul's avatar
Paul committed
100
        std::string default_value = "";
Paul's avatar
Paul committed
101
        unsigned nargs            = 1;
Paul's avatar
Paul committed
102
103
    };

Paul's avatar
Paul committed
104
    template <class T, MIGRAPHX_REQUIRES(is_multi_value<T>{})>
Paul's avatar
Paul committed
105
106
107
108
109
    std::string as_string_value(const T& x)
    {
        return to_string_range(x);
    }

kahmed10's avatar
kahmed10 committed
110
111
112
113
114
115
116
117
118
119
120
121
    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
122
    template <class T, MIGRAPHX_REQUIRES(not is_multi_value<T>{})>
Paul's avatar
Paul committed
123
124
    std::string as_string_value(const T& x)
    {
kahmed10's avatar
kahmed10 committed
125
        return as_string_value(rank<1>{}, x);
Paul's avatar
Paul committed
126
127
    }

Paul's avatar
Paul committed
128
    template <class T, class... Fs>
Paul's avatar
Paul committed
129
    void operator()(T& x, const std::vector<std::string>& flags, Fs... fs)
Paul's avatar
Paul committed
130
    {
Paul's avatar
Paul committed
131
        arguments.push_back({flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
132
133
134
135
136
                                 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
137

kahmed10's avatar
kahmed10 committed
138
139
        argument& arg = arguments.back();
        arg.type      = migraphx::get_type_name<T>();
Paul's avatar
Paul committed
140
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
kahmed10's avatar
kahmed10 committed
141
142
        if(not arg.default_value.empty() and arg.nargs > 0)
            arg.default_value = as_string_value(x);
Paul's avatar
Paul committed
143
144
    }

Paul's avatar
Paul committed
145
    template <class... Fs>
Paul's avatar
Paul committed
146
    void operator()(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
147
    {
Paul's avatar
Paul committed
148
        arguments.push_back({std::move(flags)});
Paul's avatar
Paul committed
149
150
151

        argument& arg = arguments.back();
        arg.type      = "";
Paul's avatar
Paul committed
152
        arg.nargs     = 0;
Paul's avatar
Paul committed
153
154
155
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
156
    MIGRAPHX_DRIVER_STATIC auto nargs(unsigned n = 1)
Paul's avatar
Paul committed
157
    {
Paul's avatar
Paul committed
158
        return [=](auto&&, auto& arg) { arg.nargs = n; };
Paul's avatar
Paul committed
159
160
    }

Paul's avatar
Paul committed
161
    template <class F>
Paul's avatar
Paul committed
162
    MIGRAPHX_DRIVER_STATIC auto write_action(F f)
Paul's avatar
Paul committed
163
164
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
165
            arg.action = [&, f](auto& self, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
166
167
168
169
170
171
                f(self, x, params);
                return false;
            };
        };
    }

Paul's avatar
Paul committed
172
    template <class F>
Paul's avatar
Paul committed
173
    MIGRAPHX_DRIVER_STATIC auto do_action(F f)
Paul's avatar
Paul committed
174
175
    {
        return [=](auto&, auto& arg) {
Paul's avatar
Paul committed
176
            arg.nargs  = 0;
Paul's avatar
Paul committed
177
            arg.action = [&, f](auto& self, const std::vector<std::string>&) {
Paul's avatar
Paul committed
178
179
180
181
182
183
                f(self);
                return true;
            };
        };
    }

Paul's avatar
Paul committed
184
    MIGRAPHX_DRIVER_STATIC auto append()
Paul's avatar
Paul committed
185
    {
Paul's avatar
Paul committed
186
        return write_action([](auto&, auto& x, auto& params) {
Paul's avatar
Paul committed
187
            using type = typename bare<decltype(params)>::value_type;
Paul's avatar
Paul committed
188
            std::transform(params.begin(),
Paul's avatar
Paul committed
189
190
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
191
192
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
193
194
    }

Paul's avatar
Paul committed
195
    MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
Paul's avatar
Paul committed
196
    {
Paul's avatar
Paul committed
197
        return do_action([=](auto& self) {
Paul's avatar
Paul committed
198
            for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
199
200
201
            {
                std::cout << std::endl;
                std::string prefix = "    ";
Paul's avatar
Paul committed
202
                if(arg.flags.empty())
Paul's avatar
Paul committed
203
204
205
206
                {
                    std::cout << prefix;
                    std::cout << arg.metavar;
                }
Paul's avatar
Paul committed
207
208
209
210
211
212
                for(const std::string& a : arg.flags)
                {
                    std::cout << prefix;
                    std::cout << a;
                    prefix = ", ";
                }
Paul's avatar
Paul committed
213
                if(not arg.type.empty())
Paul's avatar
Paul committed
214
                {
Paul's avatar
Paul committed
215
                    std::cout << " [" << arg.type << "]";
Paul's avatar
Paul committed
216
                    if(not arg.default_value.empty())
Paul's avatar
Paul committed
217
218
                        std::cout << " (Default: " << arg.default_value << ")";
                }
Paul's avatar
Paul committed
219
220
221
222
                std::cout << std::endl;
                std::cout << "        " << arg.help << std::endl;
            }
            std::cout << std::endl;
Paul's avatar
Paul committed
223
            if(not msg.empty())
Paul's avatar
Paul committed
224
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
225
226
227
        });
    }

Paul's avatar
Paul committed
228
    MIGRAPHX_DRIVER_STATIC auto help(const std::string& help)
Paul's avatar
Paul committed
229
    {
Paul's avatar
Paul committed
230
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
231
232
    }

Paul's avatar
Paul committed
233
    MIGRAPHX_DRIVER_STATIC auto metavar(const std::string& metavar)
Paul's avatar
Paul committed
234
235
236
237
    {
        return [=](auto&, auto& arg) { arg.metavar = metavar; };
    }

238
239
240
241
242
    MIGRAPHX_DRIVER_STATIC auto type(const std::string& type)
    {
        return [=](auto&, auto& arg) { arg.type = type; };
    }

Paul's avatar
Paul committed
243
    template <class T>
Paul's avatar
Paul committed
244
    MIGRAPHX_DRIVER_STATIC auto set_value(T value)
Paul's avatar
Paul committed
245
246
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
247
            arg.nargs  = 0;
Paul's avatar
Paul committed
248
            arg.type   = "";
Paul's avatar
Paul committed
249
            arg.action = [&, value](auto&, const std::vector<std::string>&) {
Paul's avatar
Paul committed
250
251
252
253
254
255
                x = value;
                return false;
            };
        };
    }

Paul's avatar
Paul committed
256
    bool parse(std::vector<std::string> args)
Paul's avatar
Paul committed
257
    {
Paul's avatar
Paul committed
258
        std::unordered_map<std::string, unsigned> keywords;
Paul's avatar
Paul committed
259
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
260
        {
Paul's avatar
Paul committed
261
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
262
                keywords[flag] = arg.nargs + 1;
Paul's avatar
Paul committed
263
        }
Paul's avatar
Paul committed
264
265
        auto arg_map =
            generic_parse(std::move(args), [&](const std::string& x) { return keywords[x]; });
Paul's avatar
Paul committed
266
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
267
        {
Paul's avatar
Paul committed
268
            auto flags = arg.flags;
Paul's avatar
Paul committed
269
            if(flags.empty())
Paul's avatar
Paul committed
270
                flags = {""};
Paul's avatar
Paul committed
271
            for(auto&& flag : flags)
Paul's avatar
Paul committed
272
            {
Paul's avatar
Paul committed
273
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
274
                {
Paul's avatar
Paul committed
275
                    if(arg.action(*this, arg_map[flag]))
Paul's avatar
Paul committed
276
                        return true;
Paul's avatar
Paul committed
277
278
279
                }
            }
        }
Paul's avatar
Paul committed
280
        return false;
Paul's avatar
Paul committed
281
282
283
284
285
286
287
288
289
    }

    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
290
        bool clear = false;
Paul's avatar
Paul committed
291
292
        for(auto&& x : as)
        {
Paul's avatar
Paul committed
293
294
            auto k = is_keyword(x);
            if(k > 0)
Paul's avatar
Paul committed
295
296
297
            {
                flag = x;
                result[flag]; // Ensure the flag exists
Paul's avatar
Paul committed
298
                if(k == 1)
Paul's avatar
Paul committed
299
                    flag = "";
Paul's avatar
Paul committed
300
                else if(k == 2)
Paul's avatar
Paul committed
301
302
303
                    clear = true;
                else
                    clear = false;
Paul's avatar
Paul committed
304
305
306
307
            }
            else
            {
                result[flag].push_back(x);
Paul's avatar
Paul committed
308
                if(clear)
Paul's avatar
Paul committed
309
310
                    flag = "";
                clear = false;
Paul's avatar
Paul committed
311
312
313
314
            }
        }
        return result;
    }
Paul's avatar
Paul committed
315

Paul's avatar
Paul committed
316
    private:
Paul's avatar
Paul committed
317
    std::vector<argument> arguments;
Paul's avatar
Paul committed
318
319
};

Paul's avatar
Paul committed
320
321
322
323
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
324
#endif