argument_parser.hpp 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
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
19
20
21
22
#include <migraphx/functional.hpp>

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

Paul's avatar
Paul committed
24
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
55
56
57
58
59
60
61
62
63
template <class T>
struct value_parser
{
    template <MIGRAPHX_REQUIRES(not std::is_enum<T>{})>
    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;
    }

    template <MIGRAPHX_REQUIRES(std::is_enum<T>{})>
    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);
    }
};

struct argument_parser
{
    struct argument
    {
        std::vector<std::string> flags;
Paul's avatar
Paul committed
64
        std::function<bool(argument_parser&, const std::vector<std::string>&)> action{};
Paul's avatar
Paul committed
65
66
        std::string type    = "";
        std::string help    = "";
Paul's avatar
Paul committed
67
        std::string metavar = "";
Paul's avatar
Paul committed
68
        unsigned nargs      = 1;
Paul's avatar
Paul committed
69
70
71
    };

    template <class T, class... Fs>
Paul's avatar
Paul committed
72
    void operator()(T& x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
73
    {
Paul's avatar
Paul committed
74
        arguments.push_back({flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
75
76
77
78
79
                                 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
80

Paul's avatar
Paul committed
81
82
        argument& arg = arguments.back();
        arg.type      = migraphx::get_type_name<T>();
Paul's avatar
Paul committed
83
84
85
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
86
    template <class... Fs>
Paul's avatar
Paul committed
87
    void operator()(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
Paul's avatar
Paul committed
88
    {
Paul's avatar
Paul committed
89
        arguments.push_back({std::move(flags)});
Paul's avatar
Paul committed
90
91
92

        argument& arg = arguments.back();
        arg.type      = "";
Paul's avatar
Paul committed
93
        arg.nargs     = 0;
Paul's avatar
Paul committed
94
95
96
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
97
    MIGRAPHX_DRIVER_STATIC auto nargs(unsigned n = 1)
Paul's avatar
Paul committed
98
    {
Paul's avatar
Paul committed
99
        return [=](auto&&, auto& arg) { arg.nargs = n; };
Paul's avatar
Paul committed
100
101
    }

Paul's avatar
Paul committed
102
    template <class F>
Paul's avatar
Paul committed
103
    MIGRAPHX_DRIVER_STATIC auto write_action(F f)
Paul's avatar
Paul committed
104
105
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
106
            arg.action = [&, f](auto& self, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
107
108
109
110
111
112
                f(self, x, params);
                return false;
            };
        };
    }

Paul's avatar
Paul committed
113
    template <class F>
Paul's avatar
Paul committed
114
    MIGRAPHX_DRIVER_STATIC auto do_action(F f)
Paul's avatar
Paul committed
115
116
    {
        return [=](auto&, auto& arg) {
Paul's avatar
Paul committed
117
            arg.nargs  = 0;
Paul's avatar
Paul committed
118
            arg.action = [&, f](auto& self, const std::vector<std::string>&) {
Paul's avatar
Paul committed
119
120
121
122
123
124
                f(self);
                return true;
            };
        };
    }

Paul's avatar
Paul committed
125
    MIGRAPHX_DRIVER_STATIC auto append()
Paul's avatar
Paul committed
126
    {
Paul's avatar
Paul committed
127
128
129
        return write_action([](auto&, auto& x, auto& params) {
            using type = typename decltype(params)::value_type;
            std::transform(params.begin(),
Paul's avatar
Paul committed
130
131
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
132
133
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
134
135
    }

Paul's avatar
Paul committed
136
    MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
Paul's avatar
Paul committed
137
    {
Paul's avatar
Paul committed
138
        return do_action([=](auto& self) {
Paul's avatar
Paul committed
139
            for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
140
141
142
            {
                std::cout << std::endl;
                std::string prefix = "    ";
Paul's avatar
Paul committed
143
                if(arg.flags.empty())
Paul's avatar
Paul committed
144
145
146
147
                {
                    std::cout << prefix;
                    std::cout << arg.metavar;
                }
Paul's avatar
Paul committed
148
149
150
151
152
153
                for(const std::string& a : arg.flags)
                {
                    std::cout << prefix;
                    std::cout << a;
                    prefix = ", ";
                }
Paul's avatar
Paul committed
154
                if(not arg.type.empty())
Paul's avatar
Paul committed
155
156
157
158
159
                    std::cout << " [" << arg.type << "]";
                std::cout << std::endl;
                std::cout << "        " << arg.help << std::endl;
            }
            std::cout << std::endl;
Paul's avatar
Paul committed
160
            if(not msg.empty())
Paul's avatar
Paul committed
161
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
162
163
164
        });
    }

Paul's avatar
Paul committed
165
    MIGRAPHX_DRIVER_STATIC auto help(std::string help)
Paul's avatar
Paul committed
166
    {
Paul's avatar
Paul committed
167
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
168
169
    }

Paul's avatar
Paul committed
170
    MIGRAPHX_DRIVER_STATIC auto metavar(std::string metavar)
Paul's avatar
Paul committed
171
172
173
174
    {
        return [=](auto&, auto& arg) { arg.metavar = metavar; };
    }

Paul's avatar
Paul committed
175
    template <class T>
Paul's avatar
Paul committed
176
    MIGRAPHX_DRIVER_STATIC auto set_value(T value)
Paul's avatar
Paul committed
177
178
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
179
            arg.nargs  = 0;
Paul's avatar
Paul committed
180
            arg.type   = "";
Paul's avatar
Paul committed
181
            arg.action = [&, value](auto&, const std::vector<std::string>&) {
Paul's avatar
Paul committed
182
183
184
185
186
187
                x = value;
                return false;
            };
        };
    }

Paul's avatar
Paul committed
188
    bool parse(std::vector<std::string> args)
Paul's avatar
Paul committed
189
    {
Paul's avatar
Paul committed
190
        std::unordered_map<std::string, unsigned> keywords;
Paul's avatar
Paul committed
191
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
192
        {
Paul's avatar
Paul committed
193
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
194
                keywords[flag] = arg.nargs + 1;
Paul's avatar
Paul committed
195
        }
Paul's avatar
Paul committed
196
        auto arg_map = generic_parse(std::move(args), [&](std::string x) { return keywords[x]; });
Paul's avatar
Paul committed
197
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
198
        {
Paul's avatar
Paul committed
199
            auto flags = arg.flags;
Paul's avatar
Paul committed
200
            if(flags.empty())
Paul's avatar
Paul committed
201
                flags = {""};
Paul's avatar
Paul committed
202
            for(auto&& flag : flags)
Paul's avatar
Paul committed
203
            {
Paul's avatar
Paul committed
204
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
205
                {
Paul's avatar
Paul committed
206
                    if(arg.action(*this, arg_map[flag]))
Paul's avatar
Paul committed
207
                        return true;
Paul's avatar
Paul committed
208
209
210
                }
            }
        }
Paul's avatar
Paul committed
211
        return false;
Paul's avatar
Paul committed
212
213
214
215
216
217
218
219
220
    }

    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
221
        bool clear = false;
Paul's avatar
Paul committed
222
223
        for(auto&& x : as)
        {
Paul's avatar
Paul committed
224
225
            auto k = is_keyword(x);
            if(k > 0)
Paul's avatar
Paul committed
226
227
228
            {
                flag = x;
                result[flag]; // Ensure the flag exists
Paul's avatar
Paul committed
229
                if(k == 1)
Paul's avatar
Paul committed
230
                    flag = "";
Paul's avatar
Paul committed
231
                else if(k == 2)
Paul's avatar
Paul committed
232
233
234
                    clear = true;
                else
                    clear = false;
Paul's avatar
Paul committed
235
236
237
238
            }
            else
            {
                result[flag].push_back(x);
Paul's avatar
Paul committed
239
                if(clear)
Paul's avatar
Paul committed
240
241
                    flag = "";
                clear = false;
Paul's avatar
Paul committed
242
243
244
245
            }
        }
        return result;
    }
Paul's avatar
Paul committed
246

Paul's avatar
Paul committed
247
    private:
Paul's avatar
Paul committed
248
    std::vector<argument> arguments;
Paul's avatar
Paul committed
249
250
};

Paul's avatar
Paul committed
251
252
253
254
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
255
#endif