1 files changed, 53 insertions, 35 deletions

diff --git a/includes/MappedIterator.php b/includes/MappedIterator.php
index b4376f44..70d20327 100644
--- a/includes/MappedIterator.php
+++ b/includes/MappedIterator.php
@@ -26,11 +26,15 @@
  *
  * @since 1.21
  */
-class MappedIterator implements Iterator {
-	/** @var Iterator */
-	protected $baseIterator;
-	/** @var Closure */
+class MappedIterator extends FilterIterator {
+	/** @var callable */
 	protected $vCallback;
+	/** @var callable */
+	protected $aCallback;
+	/** @var array */
+	protected $cache = array();
+
+	protected $rewound = false; // boolean; whether rewind() has been called
 
 	/**
 	 * Build an new iterator from a base iterator by having the former wrap the
@@ -38,59 +42,73 @@ class MappedIterator implements Iterator {
 	 * The callback takes the result of current() on the base iterator as an argument.
 	 * The keys of the base iterator are reused verbatim.
 	 *
+	 * An "accept" callback can also be provided which will be called for each value in
+	 * the base iterator (post-callback) and will return true if that value should be
+	 * included in iteration of the MappedIterator (otherwise it will be filtered out).
+	 *
 	 * @param Iterator|Array $iter
-	 * @param Closure $vCallback
+	 * @param callable $vCallback Value transformation callback
+	 * @param array $options Options map (includes "accept") (since 1.22)
 	 * @throws MWException
 	 */
-    public function __construct( $iter, Closure $vCallback ) {
+	public function __construct( $iter, $vCallback, array $options = array() ) {
 		if ( is_array( $iter ) ) {
-			$this->baseIterator = new ArrayIterator( $iter );
+			$baseIterator = new ArrayIterator( $iter );
 		} elseif ( $iter instanceof Iterator ) {
-			$this->baseIterator = $iter;
+			$baseIterator = $iter;
 		} else {
 			throw new MWException( "Invalid base iterator provided." );
 		}
+		parent::__construct( $baseIterator );
 		$this->vCallback = $vCallback;
-    }
+		$this->aCallback = isset( $options['accept'] ) ? $options['accept'] : null;
+	}
+
+	public function next() {
+		$this->cache = array();
+		parent::next();
+	}
 
-	/**
-	 * @return void
-	 */
 	public function rewind() {
-		$this->baseIterator->rewind();
+		$this->rewound = true;
+		$this->cache = array();
+		parent::rewind();
 	}
 
-	/**
-	 * @return Mixed|null Returns null if out of range
-	 */
-	public function current() {
-		if ( !$this->baseIterator->valid() ) {
-			return null; // out of range
+	public function accept() {
+		$value = call_user_func( $this->vCallback, $this->getInnerIterator()->current() );
+		$ok = ( $this->aCallback ) ? call_user_func( $this->aCallback, $value ) : true;
+		if ( $ok ) {
+			$this->cache['current'] = $value;
 		}
-		return call_user_func_array( $this->vCallback, array( $this->baseIterator->current() ) );
+		return $ok;
 	}
 
-	/**
-	 * @return Mixed|null Returns null if out of range
-	 */
 	public function key() {
-		if ( !$this->baseIterator->valid() ) {
-			return null; // out of range
-		}
-		return $this->baseIterator->key();
+		$this->init();
+		return parent::key();
 	}
 
-	/**
-	 * @return void
-	 */
-	public function next() {
-		$this->baseIterator->next();
+	public function valid() {
+		$this->init();
+		return parent::valid();
+	}
+
+	public function current() {
+		$this->init();
+		if ( parent::valid() ) {
+			return $this->cache['current'];
+		} else {
+			return null; // out of range
+		}
 	}
 
 	/**
-	 * @return bool
+	 * Obviate the usual need for rewind() before using a FilterIterator in a manual loop
 	 */
-	public function valid() {
-		return $this->baseIterator->valid();
+	protected function init() {
+		if ( !$this->rewound ) {
+			$this->rewind();
+		}
 	}
 }