argument_parser.hpp 8.57 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>
Paul's avatar
Paul committed
20
21
22
23

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

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

Paul's avatar
Paul committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
template<class T>
using bare = std::remove_cv_t<std::remove_reference_t<T>>;

namespace detail {

template <class T>
auto is_container(int, T&& x)
    -> decltype(x.insert(x.end(), *x.begin()), std::true_type{});

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 =
    std::integral_constant<bool,
                           (is_container<T>{} and not std::is_convertible<T, std::string>{})>;

Paul's avatar
Paul committed
55
56
57
template <class T>
struct value_parser
{
Paul's avatar
Paul committed
58
    template <MIGRAPHX_REQUIRES(not std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
59
60
61
62
63
64
65
66
67
68
69
    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
70
    template <MIGRAPHX_REQUIRES(std::is_enum<T>{} and not is_multi_value<T>{})>
Paul's avatar
Paul committed
71
72
73
74
75
76
77
78
79
80
    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
81
82
83
84
85
86
87
88
89

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

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

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

    template<class T, MIGRAPHX_REQUIRES(not is_multi_value<T>{})>
    std::string as_string_value(const T& x)
    {
        return to_string(x);
    }

Paul's avatar
Paul committed
117
    template <class T, class... Fs>
Paul's avatar
Paul committed
118
    void operator()(T& x, const std::vector<std::string>& flags, Fs... fs)
Paul's avatar
Paul committed
119
    {
Paul's avatar
Paul committed
120
        arguments.push_back({flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
121
122
123
124
125
                                 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
126

Paul's avatar
Paul committed
127
128
        argument& arg     = arguments.back();
        arg.type          = migraphx::get_type_name<T>();
Paul's avatar
Paul committed
129
        arg.default_value = as_string_value(x);
Paul's avatar
Paul committed
130
131
132
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
133
    template <class... Fs>
Paul's avatar
Paul committed
134
    void operator()(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
135
    {
Paul's avatar
Paul committed
136
        arguments.push_back({std::move(flags)});
Paul's avatar
Paul committed
137
138
139

        argument& arg = arguments.back();
        arg.type      = "";
Paul's avatar
Paul committed
140
        arg.nargs     = 0;
Paul's avatar
Paul committed
141
142
143
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
144
    MIGRAPHX_DRIVER_STATIC auto nargs(unsigned n = 1)
Paul's avatar
Paul committed
145
    {
Paul's avatar
Paul committed
146
        return [=](auto&&, auto& arg) { arg.nargs = n; };
Paul's avatar
Paul committed
147
148
    }

Paul's avatar
Paul committed
149
    template <class F>
Paul's avatar
Paul committed
150
    MIGRAPHX_DRIVER_STATIC auto write_action(F f)
Paul's avatar
Paul committed
151
152
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
153
            arg.action = [&, f](auto& self, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
154
155
156
157
158
159
                f(self, x, params);
                return false;
            };
        };
    }

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

Paul's avatar
Paul committed
172
    MIGRAPHX_DRIVER_STATIC auto append()
Paul's avatar
Paul committed
173
    {
Paul's avatar
Paul committed
174
        return write_action([](auto&, auto& x, auto& params) {
Paul's avatar
Paul committed
175
            using type = typename bare<decltype(params)>::value_type;
Paul's avatar
Paul committed
176
            std::transform(params.begin(),
Paul's avatar
Paul committed
177
178
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
179
180
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
181
182
    }

Paul's avatar
Paul committed
183
    MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
Paul's avatar
Paul committed
184
    {
Paul's avatar
Paul committed
185
        return do_action([=](auto& self) {
Paul's avatar
Paul committed
186
            for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
187
188
189
            {
                std::cout << std::endl;
                std::string prefix = "    ";
Paul's avatar
Paul committed
190
                if(arg.flags.empty())
Paul's avatar
Paul committed
191
192
193
194
                {
                    std::cout << prefix;
                    std::cout << arg.metavar;
                }
Paul's avatar
Paul committed
195
196
197
198
199
200
                for(const std::string& a : arg.flags)
                {
                    std::cout << prefix;
                    std::cout << a;
                    prefix = ", ";
                }
Paul's avatar
Paul committed
201
                if(not arg.type.empty())
Paul's avatar
Paul committed
202
                {
Paul's avatar
Paul committed
203
                    std::cout << " [" << arg.type << "]";
Paul's avatar
Paul committed
204
                    if(not arg.default_value.empty())
Paul's avatar
Paul committed
205
206
                        std::cout << " (Default: " << arg.default_value << ")";
                }
Paul's avatar
Paul committed
207
208
209
210
                std::cout << std::endl;
                std::cout << "        " << arg.help << std::endl;
            }
            std::cout << std::endl;
Paul's avatar
Paul committed
211
            if(not msg.empty())
Paul's avatar
Paul committed
212
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
213
214
215
        });
    }

Paul's avatar
Paul committed
216
    MIGRAPHX_DRIVER_STATIC auto help(const std::string& help)
Paul's avatar
Paul committed
217
    {
Paul's avatar
Paul committed
218
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
219
220
    }

Paul's avatar
Paul committed
221
    MIGRAPHX_DRIVER_STATIC auto metavar(const std::string& metavar)
Paul's avatar
Paul committed
222
223
224
225
    {
        return [=](auto&, auto& arg) { arg.metavar = metavar; };
    }

Paul's avatar
Paul committed
226
    template <class T>
Paul's avatar
Paul committed
227
    MIGRAPHX_DRIVER_STATIC auto set_value(T value)
Paul's avatar
Paul committed
228
229
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
230
            arg.nargs  = 0;
Paul's avatar
Paul committed
231
            arg.type   = "";
Paul's avatar
Paul committed
232
            arg.action = [&, value](auto&, const std::vector<std::string>&) {
Paul's avatar
Paul committed
233
234
235
236
237
238
                x = value;
                return false;
            };
        };
    }

Paul's avatar
Paul committed
239
    bool parse(std::vector<std::string> args)
Paul's avatar
Paul committed
240
    {
Paul's avatar
Paul committed
241
        std::unordered_map<std::string, unsigned> keywords;
Paul's avatar
Paul committed
242
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
243
        {
Paul's avatar
Paul committed
244
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
245
                keywords[flag] = arg.nargs + 1;
Paul's avatar
Paul committed
246
        }
Paul's avatar
Paul committed
247
248
        auto arg_map =
            generic_parse(std::move(args), [&](const std::string& x) { return keywords[x]; });
Paul's avatar
Paul committed
249
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
250
        {
Paul's avatar
Paul committed
251
            auto flags = arg.flags;
Paul's avatar
Paul committed
252
            if(flags.empty())
Paul's avatar
Paul committed
253
                flags = {""};
Paul's avatar
Paul committed
254
            for(auto&& flag : flags)
Paul's avatar
Paul committed
255
            {
Paul's avatar
Paul committed
256
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
257
                {
Paul's avatar
Paul committed
258
                    if(arg.action(*this, arg_map[flag]))
Paul's avatar
Paul committed
259
                        return true;
Paul's avatar
Paul committed
260
261
262
                }
            }
        }
Paul's avatar
Paul committed
263
        return false;
Paul's avatar
Paul committed
264
265
266
267
268
269
270
271
272
    }

    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
273
        bool clear = false;
Paul's avatar
Paul committed
274
275
        for(auto&& x : as)
        {
Paul's avatar
Paul committed
276
277
            auto k = is_keyword(x);
            if(k > 0)
Paul's avatar
Paul committed
278
279
280
            {
                flag = x;
                result[flag]; // Ensure the flag exists
Paul's avatar
Paul committed
281
                if(k == 1)
Paul's avatar
Paul committed
282
                    flag = "";
Paul's avatar
Paul committed
283
                else if(k == 2)
Paul's avatar
Paul committed
284
285
286
                    clear = true;
                else
                    clear = false;
Paul's avatar
Paul committed
287
288
289
290
            }
            else
            {
                result[flag].push_back(x);
Paul's avatar
Paul committed
291
                if(clear)
Paul's avatar
Paul committed
292
293
                    flag = "";
                clear = false;
Paul's avatar
Paul committed
294
295
296
297
            }
        }
        return result;
    }
Paul's avatar
Paul committed
298

Paul's avatar
Paul committed
299
    private:
Paul's avatar
Paul committed
300
    std::vector<argument> arguments;
Paul's avatar
Paul committed
301
302
};

Paul's avatar
Paul committed
303
304
305
306
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
307
#endif