java/pig-galago/src/com/yahoo/pig/impl/util/RewindableIterator.java

   1 /*
   2  * Copyright (c) 2007 Yahoo! Inc.  All rights reserved.
   3  * See accompanying LICENSE file.
   4  */
   5 package com.yahoo.pig.impl.util;
   6
   7 import java.io.IOException;
   8 import java.util.*;
   9
  10 public class RewindableIterator<E> {
  11     private Iterator<E> it;
  12     private ArrayList<E> buf;
  13     int pos = 0;
  14     boolean noRewind = false;
  15
  16     public RewindableIterator(Iterator<E> it) {
  17         this.it = it;
  18         buf = new ArrayList<E>();
  19     }
  20
  21     public boolean hasNext() {
  22         return (buf.size() > pos || it.hasNext());
  23     }
  24
  25     public boolean hasNext(int k) {
  26         int need = k - (buf.size() - pos);
  27
  28         while (need > 0) {
  29             if (!it.hasNext()) {
  30                 return false;
  31             } else {
  32                 buf.add(it.next());
  33                 need--;
  34             }
  35         }
  36
  37         return true;
  38     }
  39
  40     public void rewind() throws IOException {
  41         if (noRewind) throw new IOException("Internal error: attempt to rewind RewindableIterator after rewind has been disabled.");
  42         pos = 0;
  43     }
  44
  45     public void noRewind() {
  46         noRewind = true;
  47
  48         // clear out part of buffer that has been read
  49         while (pos > 0) {
  50             buf.remove(0);
  51             pos--;
  52         }
  53     }
  54
  55     public E next() {
  56         if (noRewind) {
  57             if (buf.size() <= pos) return it.next();
  58             else return buf.remove(pos);
  59         } else {
  60             if (buf.size() <= pos) buf.add(it.next());
  61
  62             return buf.get(pos++);
  63         }
  64     }
  65
  66 }