bug fix and enhancements

src/mr_util/infra/mav.h

@@ -252,6 +252,33 @@ template<typename T> class fmav: public fmav_info, public membuf<T>
     using tbuf = membuf<T>;
     using tinfo = fmav_info;
 
+    template<typename Func> void applyHelper(size_t idim, ptrdiff_t idx, Func func)
+      {
+      auto ndim = tinfo::ndim();
+      if (idim+1<ndim)
+        for (size_t i=0; i<shp[idim]; ++i)
+          applyHelper<Func>(idim+1, idx+i*str[idim], func);
+      else
+        {
+        T *d2 = vdata();
+        for (size_t i=0; i<shp[idim]; ++i)
+          func(d2[idx+i*str[idim]]);
+        }
+      }
+    template<typename Func> void applyHelper(size_t idim, ptrdiff_t idx, Func func) const
+      {
+      auto ndim = tinfo::ndim();
+      if (idim+1<ndim)
+        for (size_t i=0; i<shp[idim]; ++i)
+          applyHelper<Func>(idim+1, idx+i*str[idim], func);
+      else
+        {
+        const T *d2 = data();
+        for (size_t i=0; i<shp[idim]; ++i)
+          func(d2[idx+i*str[idim]]);
+        }
+      }
+
     auto subdata(const shape_t &i0, const shape_t &extent) const
       {
       auto ndim = tinfo::ndim();
@@ -276,7 +303,7 @@ template<typename T> class fmav: public fmav_info, public membuf<T>
           ++i2;
           }
         }
-      return make_tuple(nshp, nstr, ndim);
+      return make_tuple(nshp, nstr, nofs);
       }
 
   public:
@@ -330,20 +357,40 @@ template<typename T> class fmav: public fmav_info, public membuf<T>
       auto [nshp, nstr, nofs] = subdata(i0, extent);
       return fmav(nshp, nstr, tbuf::d+nofs, *this);
       }
+    template<typename Func> void apply(Func func)
+      {
+      if (contiguous())
+        {
+        T *d2 = vdata();
+        for (auto v=d2; v!=d2+size(); ++v)
+          func(*v);
+        return;
+        }
+      applyHelper<Func>(0, 0,func);
+      }
+    template<typename Func> void apply(Func func) const
+      {
+      if (contiguous())
+        {
+        const T *d2 = data();
+        for (auto v=d2; v!=d2+size(); ++v)
+          func(*v);
+        return;
+        }
+      applyHelper<Func>(0, 0, func);
+      }
     vector<T> dump() const
       {
-      FmavIter it(*this);
       vector<T> res(sz);
-      for (size_t i=0; i<sz; ++i, it.advance())
-        res[i] = operator[](it.ofs());
+      size_t ii=0;
+      apply([&](const T&v){res[ii++]=v;});
       return res;
       }
     void load (const vector<T> &v)
       {
       MR_assert(v.size()==sz, "bad input data size");
-      FmavIter it(*this);
-      for (size_t i=0; i<sz; ++i, it.advance())
-        vraw(it.ofs()) = v[i];
+      size_t ii=0;
+      apply([&](T &val){val=v[ii++];});
       }
   };