program BatchIterator;

{$APPTYPE CONSOLE}

uses
  Spring.Collections,
  Spring.Collections.Base;

type
  TBatchIterator<T> = class(TIterator<IEnumerable<T>>)
  private
    fSource: IEnumerable<T>;
    fSize: Integer;
    fEnumerator: IEnumerator<T>;
  protected
    function Clone: TIterator<IEnumerable<T>>; override;
  public
    constructor Create(const source: IEnumerable<T>; size: Integer);

    function MoveNext: Boolean; override;
  end;

  TEnumerableHelper = class helper for TEnumerable
    class function Batch<T>(const source: IEnumerable<T>; size: Integer): IEnumerable<IEnumerable<T>>; static;
  end;

{ TBatchIterator<T> }

constructor TBatchIterator<T>.Create(const source: IEnumerable<T>;
  size: Integer);
begin
  inherited Create;
  fSource := source;
  fSize := size;
end;

function TBatchIterator<T>.Clone: TIterator<IEnumerable<T>>;
begin
  Result := TBatchIterator<T>.Create(fSource, fSize);
end;

function TBatchIterator<T>.MoveNext: Boolean;
var
  count: Integer;
  bucket: TArray<T>;
begin
  Result := False;

  if fState = STATE_ENUMERATOR then
  begin
    fEnumerator := fSource.GetEnumerator;
    fState := STATE_RUNNING;
  end;

  if fState = STATE_RUNNING then
  begin
    count := 0;
    while (count < fSize) and fEnumerator.MoveNext do
    begin
      if not Assigned(bucket) then
        SetLength(bucket, fSize);

      bucket[count] := fEnumerator.Current;
      Inc(count);
    end;

    if Assigned(bucket) and (count > 0) then
    begin
      if count < fSize then
        SetLength(bucket, count);
      fCurrent := TEnumerable.Query<T>(bucket);
      Exit(True);
    end;
  end;
end;

{ TEnumerableHelper }

class function TEnumerableHelper.Batch<T>(const source: IEnumerable<T>;
  size: Integer): IEnumerable<IEnumerable<T>>;
begin
  Result := TBatchIterator<T>.Create(source, size);
end;

var
  batches: IEnumerable<IEnumerable<Integer>>;
  e: IEnumerable<Integer>;
  i: Integer;
begin
  batches := TEnumerable.Batch<Integer>(TEnumerable.Range(1, 1500), 1000);
  for e in batches do
  begin
    for i in e do
      Write(i, ' ');
    Writeln;
  end;

  Readln;
end.