Compositor: Add vars and methods for easier image looping

These variables and methods should make it easier to loop through buffers elements/pixels. They take into account single element buffers.
Single element buffers can be used for set operations to reduce memory usage.

Usage example: P2078

Reviewed By: #compositing, jbakker

Differential Revision: https://developer.blender.org/D11015

2 files changed, 193 insertions, 44 deletions

diff --git a/source/blender/compositor/intern/COM_MemoryBuffer.cc b/source/blender/compositor/intern/COM_MemoryBuffer.cc
index 68e39b19eaf..8c30d3215d7 100644
--- a/source/blender/compositor/intern/COM_MemoryBuffer.cc
+++ b/source/blender/compositor/intern/COM_MemoryBuffer.cc
@@ -25,29 +25,49 @@ namespace blender::compositor {
 MemoryBuffer::MemoryBuffer(MemoryProxy *memoryProxy, const rcti &rect, MemoryBufferState state)
 {
   m_rect = rect;
+  this->m_is_a_single_elem = false;
   this->m_memoryProxy = memoryProxy;
   this->m_num_channels = COM_data_type_num_channels(memoryProxy->getDataType());
   this->m_buffer = (float *)MEM_mallocN_aligned(
       sizeof(float) * buffer_len() * this->m_num_channels, 16, "COM_MemoryBuffer");
   this->m_state = state;
   this->m_datatype = memoryProxy->getDataType();
+
+  set_strides();
 }
 
-MemoryBuffer::MemoryBuffer(DataType dataType, const rcti &rect)
+MemoryBuffer::MemoryBuffer(DataType dataType, const rcti &rect, bool is_a_single_elem)
 {
   m_rect = rect;
+  this->m_is_a_single_elem = is_a_single_elem;
   this->m_memoryProxy = nullptr;
   this->m_num_channels = COM_data_type_num_channels(dataType);
   this->m_buffer = (float *)MEM_mallocN_aligned(
       sizeof(float) * buffer_len() * this->m_num_channels, 16, "COM_MemoryBuffer");
   this->m_state = MemoryBufferState::Temporary;
   this->m_datatype = dataType;
+
+  set_strides();
 }
 
 MemoryBuffer::MemoryBuffer(const MemoryBuffer &src)
-    : MemoryBuffer(src.m_memoryProxy, src.m_rect, MemoryBufferState::Temporary)
+    : MemoryBuffer(src.m_datatype, src.m_rect, false)
+{
+  m_memoryProxy = src.m_memoryProxy;
+  /* src may be single elem buffer */
+  fill_from(src);
+}
+
+void MemoryBuffer::set_strides()
 {
-  memcpy(m_buffer, src.m_buffer, buffer_len() * m_num_channels * sizeof(float));
+  if (m_is_a_single_elem) {
+    this->elem_stride = 0;
+    this->row_stride = 0;
+  }
+  else {
+    this->elem_stride = m_num_channels;
+    this->row_stride = getWidth() * m_num_channels;
+  }
 }
 
 void MemoryBuffer::clear()
@@ -100,6 +120,8 @@ MemoryBuffer::~MemoryBuffer()
 
 void MemoryBuffer::fill_from(const MemoryBuffer &src)
 {
+  BLI_assert(!this->is_a_single_elem());
+
   unsigned int otherY;
   unsigned int minX = MAX2(this->m_rect.xmin, src.m_rect.xmin);
   unsigned int maxX = MIN2(this->m_rect.xmax, src.m_rect.xmax);
@@ -109,10 +131,8 @@ void MemoryBuffer::fill_from(const MemoryBuffer &src)
   int otherOffset;
 
   for (otherY = minY; otherY < maxY; otherY++) {
-    otherOffset = ((otherY - src.m_rect.ymin) * src.getWidth() + minX - src.m_rect.xmin) *
-                  this->m_num_channels;
-    offset = ((otherY - this->m_rect.ymin) * getWidth() + minX - this->m_rect.xmin) *
-             this->m_num_channels;
+    otherOffset = src.get_coords_offset(minX, otherY);
+    offset = this->get_coords_offset(minX, otherY);
     memcpy(&this->m_buffer[offset],
            &src.m_buffer[otherOffset],
            (maxX - minX) * this->m_num_channels * sizeof(float));
@@ -123,8 +143,7 @@ void MemoryBuffer::writePixel(int x, int y, const float color[4])
 {
   if (x >= this->m_rect.xmin && x < this->m_rect.xmax && y >= this->m_rect.ymin &&
       y < this->m_rect.ymax) {
-    const int offset = (getWidth() * (y - this->m_rect.ymin) + x - this->m_rect.xmin) *
-                       this->m_num_channels;
+    const int offset = get_coords_offset(x, y);
     memcpy(&this->m_buffer[offset], color, sizeof(float) * this->m_num_channels);
   }
 }
@@ -133,8 +152,7 @@ void MemoryBuffer::addPixel(int x, int y, const float color[4])
 {
   if (x >= this->m_rect.xmin && x < this->m_rect.xmax && y >= this->m_rect.ymin &&
       y < this->m_rect.ymax) {
-    const int offset = (getWidth() * (y - this->m_rect.ymin) + x - this->m_rect.xmin) *
-                       this->m_num_channels;
+    const int offset = get_coords_offset(x, y);
     float *dst = &this->m_buffer[offset];
     const float *src = color;
     for (int i = 0; i < this->m_num_channels; i++, dst++, src++) {
@@ -151,26 +169,31 @@ static void read_ewa_pixel_sampled(void *userdata, int x, int y, float result[4]
 
 void MemoryBuffer::readEWA(float *result, const float uv[2], const float derivatives[2][2])
 {
-  BLI_assert(this->m_datatype == DataType::Color);
-  float inv_width = 1.0f / (float)this->getWidth(), inv_height = 1.0f / (float)this->getHeight();
-  /* TODO(sergey): Render pipeline uses normalized coordinates and derivatives,
-   * but compositor uses pixel space. For now let's just divide the values and
-   * switch compositor to normalized space for EWA later.
-   */
-  float uv_normal[2] = {uv[0] * inv_width, uv[1] * inv_height};
-  float du_normal[2] = {derivatives[0][0] * inv_width, derivatives[0][1] * inv_height};
-  float dv_normal[2] = {derivatives[1][0] * inv_width, derivatives[1][1] * inv_height};
-
-  BLI_ewa_filter(this->getWidth(),
-                 this->getHeight(),
-                 false,
-                 true,
-                 uv_normal,
-                 du_normal,
-                 dv_normal,
-                 read_ewa_pixel_sampled,
-                 this,
-                 result);
+  if (m_is_a_single_elem) {
+    memcpy(result, m_buffer, sizeof(float) * this->m_num_channels);
+  }
+  else {
+    BLI_assert(this->m_datatype == DataType::Color);
+    float inv_width = 1.0f / (float)this->getWidth(), inv_height = 1.0f / (float)this->getHeight();
+    /* TODO(sergey): Render pipeline uses normalized coordinates and derivatives,
+     * but compositor uses pixel space. For now let's just divide the values and
+     * switch compositor to normalized space for EWA later.
+     */
+    float uv_normal[2] = {uv[0] * inv_width, uv[1] * inv_height};
+    float du_normal[2] = {derivatives[0][0] * inv_width, derivatives[0][1] * inv_height};
+    float dv_normal[2] = {derivatives[1][0] * inv_width, derivatives[1][1] * inv_height};
+
+    BLI_ewa_filter(this->getWidth(),
+                   this->getHeight(),
+                   false,
+                   true,
+                   uv_normal,
+                   du_normal,
+                   dv_normal,
+                   read_ewa_pixel_sampled,
+                   this,
+                   result);
+  }
 }
 
 }  // namespace blender::compositor
diff --git a/source/blender/compositor/intern/COM_MemoryBuffer.h b/source/blender/compositor/intern/COM_MemoryBuffer.h
index 060a67f8797..97b220508e0 100644
--- a/source/blender/compositor/intern/COM_MemoryBuffer.h
+++ b/source/blender/compositor/intern/COM_MemoryBuffer.h
@@ -50,6 +50,25 @@ class MemoryProxy;
  * \brief a MemoryBuffer contains access to the data of a chunk
  */
 class MemoryBuffer {
+ public:
+  /**
+   * Offset between elements.
+   *
+   * Should always be used for the x dimension when calculating buffer offsets.
+   * It will be 0 when is_a_single_elem=true.
+   * e.g: buffer_index = y * buffer.row_stride + x * buffer.elem_stride
+   */
+  int elem_stride;
+
+  /**
+   * Offset between rows.
+   *
+   * Should always be used for the y dimension when calculating buffer offsets.
+   * It will be 0 when is_a_single_elem=true.
+   * e.g: buffer_index = y * buffer.row_stride + x * buffer.elem_stride
+   */
+  int row_stride;
+
  private:
   /**
    * \brief proxy of the memory (same for all chunks in the same buffer)
@@ -82,6 +101,11 @@ class MemoryBuffer {
    */
   uint8_t m_num_channels;
 
+  /**
+   * Whether buffer is a single element in memory.
+   */
+  bool m_is_a_single_elem;
+
  public:
   /**
    * \brief construct new temporarily MemoryBuffer for an area
@@ -91,7 +115,7 @@ class MemoryBuffer {
   /**
    * \brief construct new temporarily MemoryBuffer for an area
    */
-  MemoryBuffer(DataType datatype, const rcti &rect);
+  MemoryBuffer(DataType datatype, const rcti &rect, bool is_a_single_elem = false);
 
   /**
    * Copy constructor
@@ -103,6 +127,102 @@ class MemoryBuffer {
    */
   ~MemoryBuffer();
 
+  /**
+   * Whether buffer is a single element in memory independently of its resolution. True for set
+   * operations buffers.
+   */
+  bool is_a_single_elem() const
+  {
+    return m_is_a_single_elem;
+  }
+
+  float &operator[](int index)
+  {
+    BLI_assert(m_is_a_single_elem ? index < m_num_channels :
+                                    index < get_coords_offset(getWidth(), getHeight()));
+    return m_buffer[index];
+  }
+
+  const float &operator[](int index) const
+  {
+    BLI_assert(m_is_a_single_elem ? index < m_num_channels :
+                                    index < get_coords_offset(getWidth(), getHeight()));
+    return m_buffer[index];
+  }
+
+  /**
+   * Get offset needed to jump from buffer start to given coordinates.
+   */
+  int get_coords_offset(int x, int y) const
+  {
+    return (y - m_rect.ymin) * row_stride + (x - m_rect.xmin) * elem_stride;
+  }
+
+  /**
+   * Get buffer element at given coordinates.
+   */
+  float *get_elem(int x, int y)
+  {
+    BLI_assert(x >= m_rect.xmin && x < m_rect.xmax && y >= m_rect.ymin && y < m_rect.ymax);
+    return m_buffer + get_coords_offset(x, y);
+  }
+
+  /**
+   * Get buffer element at given coordinates.
+   */
+  const float *get_elem(int x, int y) const
+  {
+    BLI_assert(x >= m_rect.xmin && x < m_rect.xmax && y >= m_rect.ymin && y < m_rect.ymax);
+    return m_buffer + get_coords_offset(x, y);
+  }
+
+  /**
+   * Get channel value at given coordinates.
+   */
+  float &get_value(int x, int y, int channel)
+  {
+    BLI_assert(x >= m_rect.xmin && x < m_rect.xmax && y >= m_rect.ymin && y < m_rect.ymax &&
+               channel >= 0 && channel < m_num_channels);
+    return m_buffer[get_coords_offset(x, y) + channel];
+  }
+
+  /**
+   * Get channel value at given coordinates.
+   */
+  const float &get_value(int x, int y, int channel) const
+  {
+    BLI_assert(x >= m_rect.xmin && x < m_rect.xmax && y >= m_rect.ymin && y < m_rect.ymax &&
+               channel >= 0 && channel < m_num_channels);
+    return m_buffer[get_coords_offset(x, y) + channel];
+  }
+
+  /**
+   * Get the buffer row end.
+   */
+  const float *get_row_end(int y) const
+  {
+    BLI_assert(y >= 0 && y < getHeight());
+    return m_buffer + (is_a_single_elem() ? m_num_channels : get_coords_offset(getWidth(), y));
+  }
+
+  /**
+   * Get the number of elements in memory for a row. For single element buffers it will always
+   * be 1.
+   */
+  int get_memory_width() const
+  {
+    return is_a_single_elem() ? 1 : getWidth();
+  }
+
+  /**
+   * Get number of elements in memory for a column. For single element buffers it will
+   * always be 1.
+   */
+  int get_memory_height() const
+  {
+    return is_a_single_elem() ? 1 : getHeight();
+  }
+
   uint8_t get_num_channels()
   {
     return this->m_num_channels;
@@ -216,7 +336,7 @@ class MemoryBuffer {
       int u = x;
       int v = y;
       this->wrap_pixel(u, v, extend_x, extend_y);
-      const int offset = (getWidth() * y + x) * this->m_num_channels;
+      const int offset = get_coords_offset(u, v);
       float *buffer = &this->m_buffer[offset];
       memcpy(result, buffer, sizeof(float) * this->m_num_channels);
     }
@@ -232,7 +352,7 @@ class MemoryBuffer {
     int v = y;
 
     this->wrap_pixel(u, v, extend_x, extend_y);
-    const int offset = (getWidth() * v + u) * this->m_num_channels;
+    const int offset = get_coords_offset(u, v);
 
     BLI_assert(offset >= 0);
     BLI_assert(offset < this->buffer_len() * this->m_num_channels);
@@ -258,15 +378,20 @@ class MemoryBuffer {
       copy_vn_fl(result, this->m_num_channels, 0.0f);
       return;
     }
-    BLI_bilinear_interpolation_wrap_fl(this->m_buffer,
-                                       result,
-                                       getWidth(),
-                                       getHeight(),
-                                       this->m_num_channels,
-                                       u,
-                                       v,
-                                       extend_x == MemoryBufferExtend::Repeat,
-                                       extend_y == MemoryBufferExtend::Repeat);
+    if (m_is_a_single_elem) {
+      memcpy(result, m_buffer, sizeof(float) * this->m_num_channels);
+    }
+    else {
+      BLI_bilinear_interpolation_wrap_fl(this->m_buffer,
+                                         result,
+                                         getWidth(),
+                                         getHeight(),
+                                         this->m_num_channels,
+                                         u,
+                                         v,
+                                         extend_x == MemoryBufferExtend::Repeat,
+                                         extend_y == MemoryBufferExtend::Repeat);
+    }
   }
 
   void readEWA(float *result, const float uv[2], const float derivatives[2][2]);
@@ -321,9 +446,10 @@ class MemoryBuffer {
   float get_max_value(const rcti &rect) const;
 
  private:
+  void set_strides();
   const int buffer_len() const
   {
-    return getWidth() * getHeight();
+    return get_memory_width() * get_memory_height();
   }
 
 #ifdef WITH_CXX_GUARDEDALLOC