[python/matplotlib] Add a Dot iterator

python/port/mod/matplotlib/modpyplot.cpp

@@ -63,11 +63,11 @@ mp_obj_t modpyplot_plot(mp_obj_t x, mp_obj_t y) {
 
   // Input parameter validation
   size_t xLength, yLength;
-  mp_obj_t * xItems, yItems;
+  mp_obj_t * xItems, * yItems;
   mp_obj_get_array(x, &xLength, &xItems);
   mp_obj_get_array(y, &yLength, &yItems);
   if (xLength != yLength) {
-    mp_raise_msg_varg(&mp_type_ValueError, "x and y must have same dimension");
+    mp_raise_msg(&mp_type_ValueError, "x and y must have same dimension");
   }
 
   sPlotStore->addDots(x, y);

python/port/mod/matplotlib/plot_store.cpp

@@ -1,4 +1,5 @@
 #include "plot_store.h"
+#include <escher/palette.h>
 
 namespace Matplotlib {
 
@@ -18,49 +19,91 @@ void PlotStore::addDots(mp_obj_t x, mp_obj_t y) {
   mp_obj_list_append(m_dots, tuple);
 }
 
-PlotStore::Dot PlotStore::dotAtIndex(int i) {
+PlotStore::DotIterator PlotStore::DotIterator::Begin(mp_obj_t dots) {
+  DotIterator it;
+  mp_obj_list_get(dots, &(it.m_numberOfTuples), &(it.m_tuples));
+  if (it.m_numberOfTuples > 0) {
+    it.m_tupleIndex = 0;
+    it.loadValues();
+  }
+  return it;
+}
+
+PlotStore::DotIterator PlotStore::DotIterator::End(mp_obj_t dots) {
+  DotIterator it;
+  mp_obj_list_get(dots, &(it.m_numberOfTuples), &(it.m_tuples));
+  if (it.m_numberOfTuples > 0) {
+    it.m_tupleIndex = it.m_numberOfTuples;
+    it.m_valueIndex = 0;
+  }
+  return it;
+}
+
+PlotStore::Dot PlotStore::DotIterator::operator*() {
+  return PlotStore::Dot(
+    mp_obj_get_float(m_xValues[m_valueIndex]),
+    mp_obj_get_float(m_yValues[m_valueIndex]),
+    Palette::DataColor[m_tupleIndex] // FIXME: Share the "looping" routing
+  );
+};
+
+bool PlotStore::DotIterator::operator!=(const DotIterator & it) const {
+  return (m_tupleIndex != it.m_tupleIndex || m_valueIndex != it.m_valueIndex);
+};
+
+PlotStore::DotIterator & PlotStore::DotIterator::operator++() {
+  if (m_valueIndex < m_numberOfValues - 1) {
+    m_valueIndex++;
+  } else if (m_tupleIndex < m_numberOfTuples - 1) {
+    m_tupleIndex++;
+    loadValues();
+  } else {
+    m_tupleIndex = m_numberOfTuples;
+    m_valueIndex = 0;
+  }
+
+  return *this;
+}
+
+
+void PlotStore::DotIterator::loadValues() {
+  mp_obj_t tuple = m_tuples[m_tupleIndex];
+
+  mp_obj_t * coordinates;
+  mp_obj_get_array_fixed_n(tuple, 2, &coordinates);
+
+  mp_obj_get_array(coordinates[0], &m_numberOfValues, &m_xValues);
+  mp_obj_get_array(coordinates[1], &m_numberOfValues, &m_yValues);
+
+  m_valueIndex = 0;
+}
+
+void PlotStore::forEachDot(PlotStore::DotCallback callback) {
   size_t numberOfTuples;
   mp_obj_t * tuples;
   mp_obj_list_get(m_dots, &numberOfTuples, &tuples);
 
-  mp_obj_t firstTuple = tuples[0];
-  size_t numberOfObjects;
-  mp_obj_t * objects;
-  mp_obj_tuple_get(firstTuple, &numberOfObjects, &objects);
+  for (size_t t=0; t<numberOfTuples; t++) {
+    mp_obj_t tuple = tuples[t];
+    mp_obj_t * coordinates;
+    mp_obj_get_array_fixed_n(tuple, 2, &coordinates);
 
-  mp_obj_t x = objects[0];
-  mp_obj_t y = objects[1];
+    size_t xLength, yLength;
+    mp_obj_t * xItems, * yItems;
+    mp_obj_get_array(coordinates[0], &xLength, &xItems);
+    mp_obj_get_array(coordinates[1], &yLength, &yItems);
+    assert(xLength == yLength); // Checked in modpyplot_plot
 
-  size_t numberOfX;
-  mp_obj_t * xValues;
-  mp_obj_list_get(x, &numberOfX, &xValues);
-  size_t numberOfY;
-  mp_obj_t * yValues;
-  mp_obj_list_get(y, &numberOfY, &yValues);
+    for (size_t i=0; i<xLength; i++) {
+      PlotStore::Dot dot(
+        mp_obj_get_float(xItems[i]),
+        mp_obj_get_float(yItems[i]),
+        Palette::DataColor[t]
+      );
 
-  float goodX = mp_obj_get_float(xValues[i]);
-  float goodY = mp_obj_get_float(yValues[i]);
-
-  return PlotStore::Dot(goodX, goodY, KDColorRed);
-}
-
-int PlotStore::numberOfDots() {
-  size_t numberOfTuples;
-  mp_obj_t * tuples;
-  mp_obj_list_get(m_dots, &numberOfTuples, &tuples);
-
-  mp_obj_t firstTuple = tuples[0];
-  size_t numberOfObjects;
-  mp_obj_t * objects;
-  mp_obj_tuple_get(firstTuple, &numberOfObjects, &objects);
-
-  mp_obj_t x = objects[0];
-
-  size_t numberOfX;
-  mp_obj_t * xValues;
-  mp_obj_list_get(x, &numberOfX, &xValues);
-
-  return numberOfX;
+      callback(&dot);
+    }
+  }
 }
 
 }

python/port/mod/matplotlib/plot_store.h

@@ -26,11 +26,35 @@ public:
     KDColor m_color;
   };
 
-  // TODO: Use an iterator here. It will be a lot faster
-  Dot dotAtIndex(int i);
-  int numberOfDots();
+  class DotIterator {
+  public:
+    static DotIterator Begin(mp_obj_t dots);
+    static DotIterator End(mp_obj_t dots);
+    Dot operator*();
+    bool operator!=(const DotIterator & it) const;
+    DotIterator & operator++();
+  private:
+    void loadValues();
+    mp_obj_t * m_tuples;
+    size_t m_numberOfTuples;
+    size_t m_tupleIndex;
+    mp_obj_t * m_xValues;
+    mp_obj_t * m_yValues;
+    size_t m_numberOfValues;
+    size_t m_valueIndex;
+  };
+
+  class Dots {
+  public:
+    Dots(mp_obj_t dots) : m_dots(dots) {}
+    DotIterator begin() const { return DotIterator::Begin(m_dots); }
+    DotIterator end() const { return DotIterator::End(m_dots); }
+  private:
+    mp_obj_t m_dots;
+  };
 
   void addDots(mp_obj_t x, mp_obj_t y);
+  Dots dots() { return Dots(m_dots); }
 
   void setGrid(bool grid) { m_grid = grid; }
   bool grid() { return m_grid; }

python/port/mod/matplotlib/plot_view.cpp

@@ -11,8 +11,7 @@ void PlotView::drawRect(KDContext * ctx, KDRect rect) const {
   drawLabelsAndGraduations(ctx, rect, Axis::Vertical, true);
   drawLabelsAndGraduations(ctx, rect, Axis::Horizontal, true);
 
-  for (int i=0; i<m_store->numberOfDots(); i++) {
-    PlotStore::Dot dot = m_store->dotAtIndex(i);
+  for (PlotStore::Dot dot : m_store->dots()) {
     drawDot(ctx, rect, dot.x(), dot.y(), dot.color());
   }
 }