Shape.h

Go to the documentation of this file.

 
#ifndef KIDS_SHAPE_H
#define KIDS_SHAPE_H
 
#include <cassert>
#include <iostream>
#include <sstream>
#include <vector>
 
namespace PROJECT_NS {
 
using Basic_Dimen = std::size_t;
using Basic_Shape = std::vector<std::size_t>;
 
class Shape {
   public:
    Shape(std::size_t size) : _rank{1}, _dims{{size}}, _ldims{1}, _size{size} { enable_dynamic = false; }
 
    Shape(std::vector<std::size_t> dims) : _rank{dims.size()}, _dims{dims}, _ldims(dims.size(), 0), _size{1} {
        enable_dynamic = false;
        update();
    }
 
    Shape(std::vector<std::size_t*> dims_ptr)
        : _rank{dims_ptr.size()},
          _dims_ptr{dims_ptr},
          _dims(_dims_ptr.size(), 0),
          _ldims(_dims_ptr.size(), 0),
          _size{1} {
        enable_dynamic = true;
    }
 
    inline int rank() {
        if (enable_dynamic) update();
        return _rank;
    }
 
    inline int size() {
        if (enable_dynamic) update();
        return _size;
    }
 
    inline int operator()() { return _size; }
 
    inline std::vector<std::size_t>& dims() {
        if (enable_dynamic) update();
        return _dims;
    }
 
    inline void static_build() {
        update();
        enable_dynamic = false;
    }
 
    inline std::string to_string() {
        std::stringstream ss;
        ss << "<";
        for (auto& i : _dims) ss << i << ",";
        ss << ">";
        return ss.str();
    }
 
   private:
    void update() {
        if (enable_dynamic)
            for (int i = 0; i < _rank; ++i) {
                _dims[i] = *(_dims_ptr[i]);
                assert(_dims[i] > 0);
            }
        _ldims[_rank - 1] = 1;
        _size             = _dims[_rank - 1];
        for (int i = _rank - 2; i >= 0; --i) {
            _ldims[i] = _ldims[i + 1] * (_dims[i + 1]);
            _size *= _dims[i];
        }
    }
 
    bool enable_dynamic = false;
 
    std::size_t               _rank;      
    std::vector<std::size_t*> _dims_ptr;  
    std::vector<std::size_t>  _dims;      
    std::vector<std::size_t>  _ldims;     
    std::size_t               _size;      
};
 
};  // namespace PROJECT_NS
 
#endif  // KIDS_SHAPE_H