IVGCVSW-6059 Fixing PermutationVector.end() to cope with dimensions < 5

* PermutationVector.end() was returning the end of the fixed size array
  m_DimMappings rather than the number of mappings set by the constructor.

Signed-off-by: Colm Donelan <Colm.Donelan@arm.com>
Change-Id: Ie218f7922e8c9c35c1dc702e43a5ee2fd1a61ff0

2 files changed, 46 insertions, 1 deletions

diff --git a/include/armnn/Types.hpp b/include/armnn/Types.hpp
index 9e46d08501..de475ab68f 100644
--- a/include/armnn/Types.hpp
+++ b/include/armnn/Types.hpp
@@ -274,7 +274,11 @@ public:
     SizeType GetSize() const { return m_NumDimMappings; }
 
     ConstIterator begin() const { return m_DimMappings.begin(); }
-    ConstIterator end() const { return m_DimMappings.end(); }
+    /**
+     *
+     * @return pointer one past the end of the number of mapping not the length of m_DimMappings.
+     */
+    ConstIterator end() const { return m_DimMappings.begin() + m_NumDimMappings; }
 
     bool IsEqual(const PermutationVector& other) const
     {
diff --git a/src/armnn/test/UtilsTests.cpp b/src/armnn/test/UtilsTests.cpp
index a813feaf7f..77883ba4fb 100644
--- a/src/armnn/test/UtilsTests.cpp
+++ b/src/armnn/test/UtilsTests.cpp
@@ -269,6 +269,47 @@ BOOST_AUTO_TEST_CASE(PermuteQuantizationDim)
     BOOST_CHECK(permuted.GetQuantizationDim().value() == 3U);
 }
 
+BOOST_AUTO_TEST_CASE(PermuteVectorIterator)
+{
+    // We're slightly breaking the spirit of std::array.end() because we're using it as a
+    // variable length rather than fixed length. This test is to use a couple of iterators and
+    // make sure it still mostly makes sense.
+
+    // Create zero length.
+    armnn::PermutationVector zeroPVector({});
+    // Begin should be equal to end.
+    BOOST_CHECK(zeroPVector.begin() == zeroPVector.end());
+
+    // Create length 4. Summing the 4 values should be 6.
+    armnn::PermutationVector fourPVector({ 0, 3, 2, 1 });
+    unsigned int sum = 0;
+    for (unsigned int it : fourPVector)
+    {
+        sum += it;
+    }
+    BOOST_CHECK(sum == 6);
+    // Directly use begin and end, make sure there are 4 iterations.
+    unsigned int iterations = 0;
+    auto itr = fourPVector.begin();
+    while(itr != fourPVector.end())
+    {
+        ++iterations;
+        itr++;
+    }
+    BOOST_CHECK(iterations == 4);
+
+    // Do the same with 2 elements.
+    armnn::PermutationVector twoPVector({ 0, 1 });
+    iterations = 0;
+    itr = twoPVector.begin();
+    while(itr != twoPVector.end())
+    {
+        ++iterations;
+        itr++;
+    }
+    BOOST_CHECK(iterations == 2);
+}
+
 #if defined(ARMNNREF_ENABLED)
 BOOST_AUTO_TEST_CASE(LayerSupportHandle)
 {