argument_parser.hpp 5.66 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
24
25
26
27
28
29
30
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

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
57
        std::function<bool(argument_parser&, const std::vector<std::string>&)> action{};
Paul's avatar
Paul committed
58
59
60
61
62
63
64
65
        std::string type = "";
        std::string help = "";
    };

    template <class T, class... Fs>
    void add(T& x, std::vector<std::string> flags, Fs... fs)
    {
        arguments.emplace_back(flags, [&](auto&&, const std::vector<std::string>& params) {
Paul's avatar
Paul committed
66
            if(params.empty())
Paul's avatar
Paul committed
67
68
69
70
71
                throw std::runtime_error("Flag with no value.");
            x = value_parser<T>::apply(params.back());
            return false;
        });

Paul's avatar
Paul committed
72
73
        argument& arg = arguments.back();
        arg.type      = migraphx::get_type_name<T>();
Paul's avatar
Paul committed
74
75
76
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
77
78
79
80
81
82
83
84
85
86
    template <class... Fs>
    void add(std::nullptr_t x, std::vector<std::string> flags, Fs... fs)
    {
        arguments.push_back({flags});

        argument& arg = arguments.back();
        arg.type      = "";
        migraphx::each_args([&](auto f) { f(x, arg); }, fs...);
    }

Paul's avatar
Paul committed
87
    template <class F>
Paul's avatar
Paul committed
88
89
90
    static auto write_action(F f)
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
91
            arg.action = [&, f](auto& self, const std::vector<std::string>& params)
Paul's avatar
Paul committed
92
            {
Paul's avatar
Paul committed
93
94
95
96
97
98
                f(self, x, params);
                return false;
            };
        };
    }

Paul's avatar
Paul committed
99
    template <class F>
Paul's avatar
Paul committed
100
101
102
    static auto do_action(F f)
    {
        return [=](auto&, auto& arg) {
Paul's avatar
Paul committed
103
            arg.action = [&, f](auto& self, const std::vector<std::string>&)
Paul's avatar
Paul committed
104
            {
Paul's avatar
Paul committed
105
106
107
108
109
110
111
112
                f(self);
                return true;
            };
        };
    }

    static auto write_range()
    {
Paul's avatar
Paul committed
113
114
115
        return write_action([](auto&, auto& x, auto& params) {
            using type = typename decltype(params)::value_type;
            std::transform(params.begin(),
Paul's avatar
Paul committed
116
117
                           params.end(),
                           std::inserter(x, x.end()),
Paul's avatar
Paul committed
118
119
                           [](std::string y) { return value_parser<type>::apply(y); });
        });
Paul's avatar
Paul committed
120
121
    }

Paul's avatar
Paul committed
122
    static auto show_help(std::string msg = "")
Paul's avatar
Paul committed
123
    {
Paul's avatar
Paul committed
124
        return do_action([=](auto& self) {
Paul's avatar
Paul committed
125
            for(auto&& arg : self.arguments)
Paul's avatar
Paul committed
126
127
128
129
130
131
132
133
134
            {
                std::cout << std::endl;
                std::string prefix = "    ";
                for(const std::string& a : arg.flags)
                {
                    std::cout << prefix;
                    std::cout << a;
                    prefix = ", ";
                }
Paul's avatar
Paul committed
135
                if(not arg.type.empty())
Paul's avatar
Paul committed
136
137
138
139
140
                    std::cout << " [" << arg.type << "]";
                std::cout << std::endl;
                std::cout << "        " << arg.help << std::endl;
            }
            std::cout << std::endl;
Paul's avatar
Paul committed
141
142
            if (not msg.empty())
                std::cout << msg << std::endl;
Paul's avatar
Paul committed
143
144
145
146
147
        });
    }

    static auto help(std::string help)
    {
Paul's avatar
Paul committed
148
        return [=](auto&, auto& arg) { arg.help = help; };
Paul's avatar
Paul committed
149
150
    }

Paul's avatar
Paul committed
151
    template <class T>
Paul's avatar
Paul committed
152
153
154
    static auto set_value(T value)
    {
        return [=](auto& x, auto& arg) {
Paul's avatar
Paul committed
155
            arg.type   = "";
Paul's avatar
Paul committed
156
            arg.action = [&, value](auto&, const std::vector<std::string>&)
Paul's avatar
Paul committed
157
            {
Paul's avatar
Paul committed
158
159
160
161
162
163
164
165
166
                x = value;
                return false;
            };
        };
    }

    void parse(std::vector<std::string> args)
    {
        std::set<std::string> keywords;
Paul's avatar
Paul committed
167
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
168
169
170
        {
            keywords.insert(arg.flags.begin(), arg.flags.end());
        }
Paul's avatar
Paul committed
171
        auto arg_map = generic_parse(args, [&](std::string x) { return (keywords.count(x) > 0); });
Paul's avatar
Paul committed
172
        for(auto&& arg : arguments)
Paul's avatar
Paul committed
173
        {
Paul's avatar
Paul committed
174
            for(auto&& flag : arg.flags)
Paul's avatar
Paul committed
175
            {
Paul's avatar
Paul committed
176
                if(arg_map.count(flag) > 0)
Paul's avatar
Paul committed
177
                {
Paul's avatar
Paul committed
178
                    if(arg.action(*this, arg_map[flag]))
Paul's avatar
Paul committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
                        return;
                }
            }
        }
    }

    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;
        for(auto&& x : as)
        {
            if(is_keyword(x))
            {
                flag = x;
                result[flag]; // Ensure the flag exists
            }
            else
            {
                result[flag].push_back(x);
            }
        }
        return result;
    }
Paul's avatar
Paul committed
206

Paul's avatar
Paul committed
207
    private:
Paul's avatar
Paul committed
208
    std::vector<argument> arguments;
Paul's avatar
Paul committed
209
210
};

Paul's avatar
Paul committed
211
212
213
214
} // namespace MIGRAPHX_INLINE_NS
} // namespace driver
} // namespace migraphx

Paul's avatar
Paul committed
215
#endif