Method assign_buffer

src/argument.cpp:46–105

argument::argument(const shape& s, const argument::data_t& d) : m_shape(s), m_data(d) {}

void argument::assign_buffer(std::function<char*()> d)
{
    const shape& s = m_shape;
    if(s.type() != shape::tuple_type)
    {
        m_data = {std::move(d)};
        return;
    }
    // Collect all shapes
    std::unordered_map<std::size_t, shape> shapes;
    {
        std::size_t i = 0;
        fix([&](auto self, auto ss) {
            if(ss.sub_shapes().empty())
            {
                shapes[i] = ss;
                i++;
            }
            else
            {
                for(auto&& child : ss.sub_shapes())
                    self(child);
            }
        })(s);
    }
    // Sort by type size
    std::vector<std::size_t> order(shapes.size());
    std::iota(order.begin(), order.end(), 0);
    std::sort(order.begin(), order.end(), by(std::greater<>{}, [&](auto i) {
                  return shapes[i].type_size();
              }));
    // Compute offsets
    std::unordered_map<std::size_t, std::size_t> offsets;
    std::size_t offset = 0;
    for(auto i : order)
    {
        offsets[i] = offset;
        offset += shapes[i].bytes();
    }
    assert(offset == s.bytes());

    std::size_t i = 0;
    m_data        = fix<data_t>([&](auto self, auto ss) {
        data_t result;
        if(ss.sub_shapes().empty())
        {
            auto n = offsets[i];
            result = {[d, n]() mutable { return d() + n; }};
            i++;
            return result;
        }
        std::vector<data_t> subs;
        std::transform(ss.sub_shapes().begin(),
                       ss.sub_shapes().end(),
                       std::back_inserter(subs),
                       [&](auto child) { return self(child); });
        result.sub = subs;
        return result;