From 5ab7896fca9fa34ca0eef01e840232140188821a Mon Sep 17 00:00:00 2001
From: Teo Mrnjavac <teo@kde.org>
Date: Fri, 24 Apr 2015 16:51:34 +0200
Subject: [PATCH] Add the option of using PartitionIterator on a
 PartitionTable.

---
 .../partition/core/PartitionIterator.cpp      | 29 +++++++++++++++----
 .../partition/core/PartitionIterator.h        |  8 +++--
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/modules/partition/core/PartitionIterator.cpp b/src/modules/partition/core/PartitionIterator.cpp
index a7adc71c6..1d2112596 100644
--- a/src/modules/partition/core/PartitionIterator.cpp
+++ b/src/modules/partition/core/PartitionIterator.cpp
@@ -1,6 +1,7 @@
 /* === This file is part of Calamares - <http://github.com/calamares> ===
  *
  *   Copyright 2014, Aurélien Gâteau <agateau@kde.org>
+ *   Copyright 2015, Teo Mrnjavac <teo@kde.org>
  *
  *   Calamares is free software: you can redistribute it and/or modify
  *   it under the terms of the GNU General Public License as published by
@@ -23,8 +24,8 @@
 #include <core/partition.h>
 #include <core/partitiontable.h>
 
-PartitionIterator::PartitionIterator( Device* device )
-    : m_device( device )
+PartitionIterator::PartitionIterator( PartitionTable* table )
+    : m_table( table )
 {}
 
 Partition*
@@ -70,7 +71,7 @@ PartitionIterator::operator++()
 bool
 PartitionIterator::operator==( const PartitionIterator& other ) const
 {
-    return m_device == other.m_device && m_current == other.m_current;
+    return m_table == other.m_table && m_current == other.m_current;
 }
 
 bool
@@ -82,10 +83,16 @@ PartitionIterator::operator!=( const PartitionIterator& other ) const
 PartitionIterator
 PartitionIterator::begin( Device* device )
 {
-    auto it = PartitionIterator( device );
     PartitionTable* table = device->partitionTable();
     if ( !table )
-        return it;
+        return PartitionIterator( nullptr );
+    return PartitionIterator::begin( table );
+}
+
+PartitionIterator
+PartitionIterator::begin( PartitionTable* table )
+{
+    auto it = PartitionIterator( table );
     QList< Partition* > children = table->children();
     // Does not usually happen, but it did happen on a 10MB disk with an MBR
     // partition table.
@@ -98,5 +105,15 @@ PartitionIterator::begin( Device* device )
 PartitionIterator
 PartitionIterator::end( Device* device )
 {
-    return PartitionIterator( device );
+    PartitionTable* table = device->partitionTable();
+    if ( !table )
+        return PartitionIterator( nullptr );
+
+    return PartitionIterator::end( table );
+}
+
+PartitionIterator
+PartitionIterator::end( PartitionTable* table )
+{
+    return PartitionIterator( table );
 }
diff --git a/src/modules/partition/core/PartitionIterator.h b/src/modules/partition/core/PartitionIterator.h
index 7b527b241..225022273 100644
--- a/src/modules/partition/core/PartitionIterator.h
+++ b/src/modules/partition/core/PartitionIterator.h
@@ -1,6 +1,7 @@
 /* === This file is part of Calamares - <http://github.com/calamares> ===
  *
  *   Copyright 2014, Aurélien Gâteau <agateau@kde.org>
+ *   Copyright 2015, Teo Mrnjavac <teo@kde.org>
  *
  *   Calamares is free software: you can redistribute it and/or modify
  *   it under the terms of the GNU General Public License as published by
@@ -21,6 +22,7 @@
 
 class Device;
 class Partition;
+class PartitionTable;
 
 /**
  * A forward-only iterator to go through the partitions of a device,
@@ -37,12 +39,14 @@ public:
     bool operator!=( const PartitionIterator& other ) const;
 
     static PartitionIterator begin( Device* device );
+    static PartitionIterator begin( PartitionTable* table );
     static PartitionIterator end( Device* device );
+    static PartitionIterator end( PartitionTable* table );
 
 private:
-    PartitionIterator( Device* device );
+    PartitionIterator( PartitionTable* table );
 
-    Device* m_device;
+    PartitionTable* m_table;
     Partition* m_current = nullptr;
 };