From 32e0026128d2566d257bc3122cdbdabeaafec445 Mon Sep 17 00:00:00 2001
From: Matthias Hochsteger <mhochsteger@cerbsim.com>
Date: Mon, 2 Dec 2024 19:35:49 +0100
Subject: [PATCH] Some memory tracer fixes/features

If range checks are enabled:
- Trace all objects
- Check if memory usage never gets negative
- Check if memory usage is 0 in destructor
- Track total memory usage (use pyngcore.GetTotalMemory())
---
 libsrc/core/array.hpp                |  26 ++++---
 libsrc/core/bitarray.cpp             |   5 +-
 libsrc/core/bitarray.hpp             |   9 ++-
 libsrc/core/exception.cpp            |   6 +-
 libsrc/core/exception.hpp            |  19 ++---
 libsrc/core/memtracer.hpp            | 102 ++++++++++++++++++++-------
 libsrc/core/profiler.cpp             |   1 +
 libsrc/core/python_ngcore_export.cpp |   2 +
 libsrc/core/table.hpp                |   9 ++-
 9 files changed, 126 insertions(+), 53 deletions(-)

diff --git a/libsrc/core/array.hpp b/libsrc/core/array.hpp
index 51ed0d10..462f56b2 100644
--- a/libsrc/core/array.hpp
+++ b/libsrc/core/array.hpp
@@ -687,6 +687,7 @@ namespace ngcore
     size_t allocsize;
     /// that's the data we have to delete, nullptr for not owning the memory
     T * mem_to_delete;
+    MemoryTracer mt;
 
 
     using FlatArray<T,IndexType>::size;
@@ -708,6 +709,7 @@ namespace ngcore
     {
       allocsize = asize; 
       mem_to_delete = data;
+      mt.Alloc(sizeof(T)*asize);
     }
 
 
@@ -717,7 +719,10 @@ namespace ngcore
     {
       allocsize = asize;
       if(ownMemory)
+      {
         mem_to_delete = adata;
+        mt.Alloc(sizeof(T)*asize);
+      }
       else
         mem_to_delete = nullptr;
     }
@@ -733,8 +738,7 @@ namespace ngcore
 
     NETGEN_INLINE Array (Array && a2) 
     {
-      mt.Swap(0., a2.mt, sizeof(T) * a2.allocsize);
-
+      mt = std::move(a2.mt);
       size = a2.size; 
       data = a2.data;
       allocsize = a2.allocsize;
@@ -753,6 +757,7 @@ namespace ngcore
         {
           allocsize = size;
           mem_to_delete = data;
+          mt.Alloc(sizeof(T)*size);
           for (size_t i = 0; i < size; i++)
             data[i] = a2.data[i];
         }
@@ -772,6 +777,7 @@ namespace ngcore
     {
       allocsize = size;
       mem_to_delete = data;
+      mt.Alloc(sizeof(T)*size);
       /*
       for (size_t i = 0; i < size; i++)
         data[i] = a2[i];
@@ -788,6 +794,7 @@ namespace ngcore
     {
       allocsize = size;
       mem_to_delete = data;
+      mt.Alloc(sizeof(T)*size);
       size_t cnt = 0;
       for (auto val : list)
         data[cnt++] = val;
@@ -800,6 +807,7 @@ namespace ngcore
     {
       allocsize = size;
       mem_to_delete = data;
+      mt.Alloc(sizeof(T)*size);
       for(size_t i = 0; i <  a2.Size(); i++)
         data[i] = a2[i];
       for (size_t i = a2.Size(), j=0; i < size; i++,j++)
@@ -1011,8 +1019,7 @@ namespace ngcore
     /// steal array 
     NETGEN_INLINE Array & operator= (Array && a2)
     {
-      mt.Swap(sizeof(T)*allocsize, a2.mt, sizeof(T)*a2.allocsize);
-
+      mt = std::move(a2.mt);
       ngcore::Swap (size, a2.size);
       ngcore::Swap (data, a2.data);
       ngcore::Swap (allocsize, a2.allocsize);
@@ -1086,8 +1093,7 @@ namespace ngcore
     
     NETGEN_INLINE void Swap (Array & b)
     {
-      mt.Swap(sizeof(T) * allocsize, b.mt, sizeof(T) * b.allocsize);
-
+      mt = std::move(b.mt);
       ngcore::Swap (size, b.size);
       ngcore::Swap (data, b.data);
       ngcore::Swap (allocsize, b.allocsize);
@@ -1096,7 +1102,8 @@ namespace ngcore
 
     NETGEN_INLINE void StartMemoryTracing () const
     {
-      mt.Alloc(sizeof(T) * allocsize);
+      if(mem_to_delete)
+        mt.Alloc(sizeof(T) * allocsize);
     }
 
     const MemoryTracer& GetMemoryTracer() const { return mt; }
@@ -1105,7 +1112,6 @@ namespace ngcore
 
     /// resize array, at least to size minsize. copy contents
     NETGEN_INLINE void ReSize (size_t minsize);
-    MemoryTracer mt;
   };
 
   
@@ -1158,6 +1164,7 @@ namespace ngcore
     using Array<T>::allocsize;
     using Array<T>::data;
     using Array<T>::mem_to_delete;
+    using Array<T>::mt;
     // using Array<T>::ownmem;
 
   public:
@@ -1171,6 +1178,7 @@ namespace ngcore
           data = new T[asize];
           allocsize = size;
           mem_to_delete = data;
+          mt.Alloc(sizeof(T)*asize);
         }
     }
 
@@ -1191,6 +1199,7 @@ namespace ngcore
     ArrayMem(ArrayMem && a2)
       : Array<T> (a2.Size(), (T*)mem)
     {
+      mt = std::move(a2.mt);
       if (a2.mem_to_delete)
         {
           mem_to_delete = a2.mem_to_delete;
@@ -1233,6 +1242,7 @@ namespace ngcore
 
     ArrayMem & operator= (ArrayMem && a2)
     {
+      mt = std::move(a2.mt);
       ngcore::Swap (mem_to_delete, a2.mem_to_delete);
       ngcore::Swap (allocsize, a2.allocsize);
       ngcore::Swap (size, a2.size);
diff --git a/libsrc/core/bitarray.cpp b/libsrc/core/bitarray.cpp
index 8c76ba90..1c6f6ec4 100644
--- a/libsrc/core/bitarray.cpp
+++ b/libsrc/core/bitarray.cpp
@@ -40,12 +40,13 @@ namespace ngcore
     if (owns_data)
       {
         delete [] data;
-        mt.Free(Addr(size)+1);
+        mt.Free(GetMemoryUsage());
       }
 
     size = asize;
     data = new unsigned char [Addr (size)+1];
-    mt.Alloc(Addr(size)+1);
+    owns_data = true;
+    mt.Alloc(GetMemoryUsage());
   }
 
   BitArray & BitArray :: Set () throw()
diff --git a/libsrc/core/bitarray.hpp b/libsrc/core/bitarray.hpp
index 8ae32e3c..91259e91 100644
--- a/libsrc/core/bitarray.hpp
+++ b/libsrc/core/bitarray.hpp
@@ -49,6 +49,7 @@ public:
   {
     ba2.owns_data = false;
     ba2.data = nullptr;
+    mt = std::move(ba2.mt);
   }
 
   template <typename T>
@@ -59,13 +60,17 @@ public:
     int cnt = 0;
     for (auto i = list.begin(); i < list.end(); i++, cnt++)
       if (*i) SetBit(cnt);
+    StartMemoryTracing();
   }
 
   /// delete data
   ~BitArray ()
   {
     if (owns_data)
+    {
       delete [] data;
+      mt.Free(GetMemoryUsage());
+    }
   }
 
   /// Set size, loose values
@@ -150,11 +155,11 @@ public:
   
   NGCORE_API auto * Data() const { return data; }
 
+  const size_t GetMemoryUsage() const { return owns_data ? (size+CHAR_BIT-1)/CHAR_BIT : 0; }
   const MemoryTracer& GetMemoryTracer() const { return mt; }
   void StartMemoryTracing() const
   {
-    if(owns_data)
-      mt.Alloc(Addr(size)+1);
+    mt.Alloc(GetMemoryUsage());
   }
 
 private:
diff --git a/libsrc/core/exception.cpp b/libsrc/core/exception.cpp
index 8e2f251e..1411b25e 100644
--- a/libsrc/core/exception.cpp
+++ b/libsrc/core/exception.cpp
@@ -50,7 +50,7 @@ namespace ngcore
   
   RangeException :: RangeException (// const std::string & where,
                                     const char * where,
-                                    int ind, int imin, int imax) : Exception("")
+                                    ptrdiff_t ind, ptrdiff_t imin, ptrdiff_t imax) : Exception("")
   {
     std::stringstream str;
     str << where << ": index " << ind << " out of range [" << imin << "," << imax << ")\n";
@@ -59,7 +59,7 @@ namespace ngcore
   }
   
 
-  void ThrowRangeException(const char * s, int ind, int imin, int imax)
+  void ThrowRangeException(const char * s, ptrdiff_t ind, ptrdiff_t imin, ptrdiff_t imax)
   {
     throw RangeException(s, ind, imin, imax);
   }
@@ -75,7 +75,7 @@ namespace ngcore
   }
 
 
-  void ThrowNotTheSameException(const char * s, long int a, long int b)
+  void ThrowNotTheSameException(const char * s, ptrdiff_t a, ptrdiff_t b)
   {
     throw ngcore::Exception(std::string(s) + ", a="+ToString(a) + ", b="+ToString(b) + GetBackTrace());     
   }
diff --git a/libsrc/core/exception.hpp b/libsrc/core/exception.hpp
index 72d0ce9c..8615afc1 100644
--- a/libsrc/core/exception.hpp
+++ b/libsrc/core/exception.hpp
@@ -1,6 +1,7 @@
 #ifndef NETGEN_CORE_EXCEPTION_HPP
 #define NETGEN_CORE_EXCEPTION_HPP
 
+#include <cstddef>
 #include <sstream>         // for stringstream
 #include <stdexcept>       // for exception
 #include <string>          // for string
@@ -66,7 +67,7 @@ namespace ngcore
     /// where it occurs, index, minimal and maximal indices
     RangeException (// const std::string & where,
                     const char * where,
-                    int ind, int imin, int imax);
+                    ptrdiff_t ind, ptrdiff_t imin, ptrdiff_t imax);
     /*
     : Exception("")
       {
@@ -85,8 +86,8 @@ namespace ngcore
     }
   };
 
-  [[noreturn]] NGCORE_API void ThrowRangeException(const char * s, int ind, int imin, int imax);
-  [[noreturn]] NGCORE_API void ThrowNotTheSameException(const char * s, long int a, long int b);
+  [[noreturn]] NGCORE_API void ThrowRangeException(const char * s, ptrdiff_t ind, ptrdiff_t imin, ptrdiff_t imax);
+  [[noreturn]] NGCORE_API void ThrowNotTheSameException(const char * s, ptrdiff_t a, ptrdiff_t b);
   
   
   // Exception used if no simd implementation is available to fall back to standard evaluation
@@ -98,12 +99,12 @@ namespace ngcore
     constexpr operator bool() const { return false; } };
 
   namespace detail {
-    template <typename T>
-    inline static void CheckRange(const char * s, const T& n, int first, int next)
+    template <typename T, typename Tmin, typename Tmax>
+    inline static void CheckRange(const char * s, const T& n, Tmin first, Tmax next)
     {
       if constexpr (!IsSafe<decltype(n)>())
         if (n<first || n>=next)
-          ThrowRangeException(s, int(n), first, next);
+          ThrowRangeException(s, ptrdiff_t(n), ptrdiff_t(first), ptrdiff_t(next));
     }
 
     template <typename Ta, typename Tb>
@@ -112,10 +113,10 @@ namespace ngcore
      if constexpr (!IsSafe<decltype(a)>() || !IsSafe<decltype(b)>())
       if(a != b)
       {
-        if constexpr(std::is_same<decltype(a),size_t>() && std::is_same<decltype(b),size_t>())
+        if constexpr(std::is_integral_v<decltype(a)> && std::is_same_v<decltype(a),decltype(b)>)
           ThrowNotTheSameException(s, long(a), long(b)); \
         else
-          throw Exception(std::string(s) + "\t: not the same"+ToString(a) + ", b="+ngcore::ToString(b) + GetBackTrace());
+          throw Exception(std::string(s) + "\t: not the same, a="+ToString(a) + ", b="+ngcore::ToString(b) + GetBackTrace());
       }
     }
   } // namespace detail
@@ -128,7 +129,7 @@ namespace ngcore
 #define NG_EXCEPTION(s) ngcore::Exception(__FILE__ ":" NETGEN_CORE_NGEXEPTION_STR(__LINE__) "\t"+std::string(s))
 
 #if defined(NETGEN_ENABLE_CHECK_RANGE) && !defined(__CUDA_ARCH__)
-#define NETGEN_CHECK_RANGE(value, min, max_plus_one) ngcore::detail::CheckRange(__FILE__ ":" NETGEN_CORE_NGEXEPTION_STR(__LINE__) "\t", value, int(min), int(max_plus_one));
+#define NETGEN_CHECK_RANGE(value, min, max_plus_one) ngcore::detail::CheckRange(__FILE__ ":" NETGEN_CORE_NGEXEPTION_STR(__LINE__) "\t", value, min, max_plus_one);
 #define NETGEN_CHECK_SAME(a,b) ngcore::detail::CheckSame(__FILE__ ":" NETGEN_CORE_NGEXEPTION_STR(__LINE__) "\t", a, b);
 
 #define NETGEN_NOEXCEPT 
diff --git a/libsrc/core/memtracer.hpp b/libsrc/core/memtracer.hpp
index 2e5a0e30..757e6405 100644
--- a/libsrc/core/memtracer.hpp
+++ b/libsrc/core/memtracer.hpp
@@ -35,11 +35,16 @@ namespace ngcore
 
   class MemoryTracer
   {
-    #if defined(NETGEN_TRACE_MEMORY) && !defined(__CUDA_ARCH__)
+#if defined(NETGEN_TRACE_MEMORY) && !defined(__CUDA_ARCH__)
     NGCORE_API static std::vector<std::string> names;
     NGCORE_API static std::vector<int> parents;
 
-    static int CreateId(const std::string& name)
+    #if defined(NETGEN_CHECK_RANGE)
+    NGCORE_API static std::atomic<size_t> total_memory;
+    mutable size_t allocated_memory = 0;
+    #endif // NETGEN_CHECK_RANGE
+
+    static int CreateId(const std::string& name = "")
     {
       int id = names.size();
       names.push_back(name);
@@ -48,7 +53,7 @@ namespace ngcore
         std::cerr << "Allocated " << id << " MemoryTracer objects" << std::endl;
       return id;
     }
-    int id;
+    mutable int id = 0;
 
     public:
 
@@ -57,8 +62,33 @@ namespace ngcore
       id = CreateId(name);
     }
 
-    // not tracing
-    MemoryTracer() : id(0) {}
+    MemoryTracer() { }
+
+    MemoryTracer(const MemoryTracer & tracer)
+    {
+      (*this) = tracer;
+    }
+
+    MemoryTracer(MemoryTracer && tracer)
+    {
+      (*this) = std::move(tracer);
+    }
+
+    MemoryTracer & operator=(const MemoryTracer & tracer) {
+      if(tracer.id)
+        id = CreateId(names[tracer.id]);
+      return *this;
+    }
+
+    MemoryTracer & operator=(MemoryTracer && tracer) {
+      ngcore::Swap(id, tracer.id);
+
+      #if defined(NETGEN_CHECK_RANGE)
+      ngcore::Swap(allocated_memory, tracer.allocated_memory);
+      #endif // NETGEN_CHECK_RANGE
+
+      return *this;
+    }
 
     template <typename... TRest>
     MemoryTracer( std::string name, TRest & ... rest )
@@ -67,38 +97,48 @@ namespace ngcore
       Track(rest...);
     }
 
+    #if defined(NETGEN_CHECK_RANGE)
+    // check if all memory was freed when object is destroyed
+    ~MemoryTracer()
+    {
+      NETGEN_CHECK_SAME(allocated_memory, 0);
+    }
+    #endif // NETGEN_CHECK_RANGE
+
     NETGEN_INLINE void Alloc(size_t size) const
     {
+      #if defined(NETGEN_CHECK_RANGE)
+      // Trace also nameless Memtracer objects if range checks are active
+      if(!id && size)
+        id = CreateId();
+      #endif // NETGEN_CHECK_RANGE
+
       if(id && trace)
         trace->AllocMemory(id, size);
+
+      #if defined(NETGEN_CHECK_RANGE)
+      if(id)
+      {
+        allocated_memory += size;
+        total_memory += size;
+      }
+      #endif // NETGEN_CHECK_RANGE
     }
 
     void Free(size_t size) const
     {
       if(id && trace)
         trace->FreeMemory(id, size);
-    }
 
-    void Swap(size_t mysize, MemoryTracer& other, size_t other_size) const
-    {
-      if(!trace || (id == 0 && other.id == 0))
-        return;
-      if(id == 0)
-        return trace->ChangeMemory(other.id, mysize - other_size);
-      if(other.id == 0)
-        return trace->ChangeMemory(id, other_size - mysize);
-
-      // first decrease memory, otherwise have artificial/wrong high peak memory usage
-      if(mysize<other_size)
-        {
-          trace->ChangeMemory(other.id, mysize-other_size);
-          trace->ChangeMemory(id, other_size-mysize);
-        }
-      else
-        {
-          trace->ChangeMemory(id, other_size-mysize);
-          trace->ChangeMemory(other.id, mysize-other_size);
-        }
+      #if defined(NETGEN_CHECK_RANGE)
+      if(id)
+      {
+        // check if we have at least size bytes of memory currently allocated (such that allocated_memory doesn't get negative)
+        NETGEN_CHECK_RANGE(allocated_memory, static_cast<ptrdiff_t>(size), std::numeric_limits<ptrdiff_t>::max());
+        allocated_memory -= size;
+        total_memory -= size;
+        #endif // NETGEN_CHECK_RANGE
+      }
     }
 
     int GetId() const { return id; }
@@ -148,6 +188,14 @@ namespace ngcore
 
     static const std::vector<std::string> & GetNames() { return names; }
     static const std::vector<int> & GetParents() { return parents; }
+    static size_t GetTotalMemory()
+    {
+      #if defined(NETGEN_CHECK_RANGE)
+      return total_memory;
+      #else
+      return 0;
+      #endif // NETGEN_CHECK_RANGE
+    }
 #else // defined(NETGEN_TRACE_MEMORY) && !defined(__CUDA_ARCH__)
   public:
     MemoryTracer() {}
@@ -157,7 +205,6 @@ namespace ngcore
 
     void Alloc(size_t /* size */) const {}
     void Free(size_t /* size */) const {}
-    void Swap(...) const {}
     int GetId() const { return 0; }
 
     template <typename... TRest>
@@ -166,6 +213,7 @@ namespace ngcore
     static std::string GetName(int /* id */) { return ""; }
     std::string GetName() const { return ""; }
     void SetName(std::string /* name */) const {}
+    static size_t GetTotalMemory() { return 0; }
 #endif // NETGEN_TRACE_MEMORY
   };
 } // namespace ngcore
diff --git a/libsrc/core/profiler.cpp b/libsrc/core/profiler.cpp
index cf49654b..6f257cfb 100644
--- a/libsrc/core/profiler.cpp
+++ b/libsrc/core/profiler.cpp
@@ -116,6 +116,7 @@ namespace ngcore
 #ifdef NETGEN_TRACE_MEMORY
   std::vector<std::string> MemoryTracer::names{"all"};
   std::vector<int> MemoryTracer::parents{-1};
+  std::atomic<size_t> MemoryTracer::total_memory{0};
 #endif // NETGEN_TRACE_MEMORY
 
 } // namespace ngcore
diff --git a/libsrc/core/python_ngcore_export.cpp b/libsrc/core/python_ngcore_export.cpp
index eea84619..cc9f5a70 100644
--- a/libsrc/core/python_ngcore_export.cpp
+++ b/libsrc/core/python_ngcore_export.cpp
@@ -324,6 +324,8 @@ threads : int
 #endif // NETGEN_TRACE_MEMORY
     ;
 
+    m.def("GetTotalMemory", MemoryTracer::GetTotalMemory);
+
     py::class_<Timer<>> (m, "Timer")
     .def(py::init<const string&>())
     .def("Start", static_cast<void (Timer<>::*)()const>(&Timer<>::Start), "start timer")
diff --git a/libsrc/core/table.hpp b/libsrc/core/table.hpp
index 47893ed5..134ed8a9 100644
--- a/libsrc/core/table.hpp
+++ b/libsrc/core/table.hpp
@@ -130,6 +130,7 @@ namespace ngcore
     {
       for (size_t i : IntRange(size+1))
         index[i] = i*entrysize;
+      mt.Alloc(GetMemUsage());
     }
 
     /// Construct table of variable entrysize
@@ -141,6 +142,7 @@ namespace ngcore
       index = TablePrefixSum (FlatArray<TI> (entrysize.Size(), entrysize.Data()));
       size_t cnt = index[size];
       data = new T[cnt];
+      mt.Alloc(GetMemUsage());
     }
 
     explicit NETGEN_INLINE Table (const FlatTable<T,IndexType> & tab2)
@@ -157,6 +159,7 @@ namespace ngcore
       size_t cnt = index[size];
       data = new T[cnt];
       this->AsArray() = tab2.AsArray();
+      mt.Alloc(GetMemUsage());
       /*
       for (size_t i = 0; i < cnt; i++)
         data[i] = tab2.data[i];
@@ -177,12 +180,14 @@ namespace ngcore
       data = new T[cnt];
       for (size_t i = 0; i < cnt; i++)
         data[i] = tab2.data[i];
+
+      mt.Alloc(GetMemUsage());
     }
 
     NETGEN_INLINE Table (Table && tab2)
       : FlatTable<T,IndexType>(0, nullptr, nullptr)
     {
-      tab2.mt.Free(tab2.GetMemUsage());
+      mt = std::move(tab2.mt);
       Swap (size, tab2.size);
       Swap (index, tab2.index);
       Swap (data, tab2.data);
@@ -210,7 +215,7 @@ namespace ngcore
 
     NETGEN_INLINE Table & operator= (Table && tab2)
     {
-      mt.Swap(GetMemUsage(), tab2.mt, tab2.GetMemUsage());
+      mt = std::move(tab2.mt);
       Swap (size, tab2.size);
       Swap (index, tab2.index);
       Swap (data, tab2.data);