Iter.mag

declare type PGGIter;
declare attributes PGGIter: cache;

declare type PGGIter_List: PGGIter;
declare attributes PGGIter_List: list, idx;

declare type PGGIter_Filter: PGGIter;
declare attributes PGGIter_Filter: iter, filter;

declare type PGGIter_Truncate: PGGIter;
declare attributes PGGIter_Truncate: iter, limit;

declare type PGGIter_Enumerate: PGGIter;
declare attributes PGGIter_Enumerate: iter, count;

declare type PGGIter_Apply: PGGIter;
declare attributes PGGIter_Apply: iter, map;

declare type PGGIter_User: PGGIter;
declare attributes PGGIter_User: state, has_next;

intrinsic PGG_Iter(has_next :: UserProgram : State:=false) -> PGGIter
  {Makes an iterator X with whose state attribute is State and `HasNext(X)` is `has_next(X)`.}
  X := New(PGGIter_User);
  X`state := State;
  X`has_next := has_next;
  return X;
end intrinsic;

intrinsic HasNext(X :: PGGIter_User) -> BoolElt, .
  {True if X has a next item.}
  return X`has_next(X);
end intrinsic;

intrinsic PGG_ToIter(xs :: []) -> PGGIter
  {Makes an iterator.}
  X := New(PGGIter_List);
  X`list := SequenceToList(xs);
  X`idx := 0;
  return X;
end intrinsic;

intrinsic PGG_ToIter(xs :: List) -> PGGIter
  {Makes an iterator.}
  X := New(PGGIter_List);
  X`list := xs;
  X`idx := 0;
  return X;
end intrinsic;

intrinsic HasNext(X :: PGGIter_List) -> BoolElt, .
  {True if X has a next item.}
  if X`idx lt #X`list then
    X`idx +:= 1;
    return true, X`list[X`idx];
  else
    return false, _;
  end if;
end intrinsic;

intrinsic ToList(X :: PGGIter) -> List
  {Converts X to a list.}
  ret := [];
  while true do
    ok, x := HasNext(X);
    if ok then
      Append(~ret, x);
    else
      return ret;
    end if;
  end while;
end intrinsic;

// we make a specialized version for the common case of just dealing with lists directly
intrinsic ToList(X :: PGGIter_List) -> List
  {"}
  if X`idx eq 0 then
    return X`list;
  else
    return X`list[X`idx+1..#X`list];
  end if;
end intrinsic;

intrinsic ToSequence(X :: PGGIter) -> []
  {Converts X to a sequence.}
  return [x : x in ToList(X)];
end intrinsic;

intrinsic Enumerate(X :: PGGIter) -> PGGIter
  {The sequence <i,x_i> for x_i in X.}
  Y := New(PGGIter_Enumerate);
  Y`iter := X;
  Y`count := 0;
  return Y;
end intrinsic;

intrinsic HasNext(X :: PGGIter_Enumerate) -> BoolElt, .
  {True if X has a next item.}
  ok, x := HasNext(X`iter);
  if ok then
    X`count +:= 1;
    return true, <X`count, x>;
  else
    return false, _;
  end if;
end intrinsic;

intrinsic Apply(X :: PGGIter, f) -> PGGIter
  {The sequence f(x) for x in X.}
  Y := New(PGGIter_Apply);
  Y`iter := X;
  Y`map := f;
  return Y;
end intrinsic;

intrinsic HasNext(X :: PGGIter) -> BoolElt, .
  {True if X has a next item.}
  ok, x := HasNext(X`iter);
  if ok then
    return true, X`map(x);
  else
    return false, _;
  end if;
end intrinsic;

intrinsic Filter(X :: PGGIter, f) -> PGGIter
  {The x in X such that f(x) is true.}
  Y := New(PGGIter_Filter);
  Y`iter := X;
  Y`filter := f;
  return Y;
end intrinsic;

intrinsic HasNext(X :: PGGIter_Filter) -> BoolElt, .
  {True if X has a next item.}
  idx := 0;
  while true do
    ok, x := HasNext(X`iter);
    if ok then
      if X`filter(x) then
        return true, x;
      end if;
    else
      return false, _;
    end if;
  end while;
end intrinsic;

intrinsic Reverse(X :: PGGIter) -> PGGIter
  {The reverse.}
  return PGG_ToIter(Reverse(ToList(X)));
end intrinsic;

intrinsic Shuffle(X :: PGGIter) -> PGGIter
  {Randomize the order of X.}
  xs := ToList(X);
  if #xs eq 0 then
    return X;
  else
    g := Random(SymmetricGroup(#xs));
    return PGG_ToIter([* xs[i^g] : i in [1..#xs] *]);
  end if;
end intrinsic;

intrinsic SortBy(X :: PGGIter, keyfunc :: UserProgram) -> PGGIter
  {Sorts X so that keyfunc(x) are naturally sorted.}
  xs := ToList(X);
  keys := [keyfunc(x) : x in xs];
  Sort(~keys, ~permut);
  return PGG_ToIter([* xs[i^permut] : i in [1..#xs] *]);
end intrinsic;

intrinsic FilterHasNext(X :: PGGIter, f) -> BoolElt, ., RngIntElt
  {True if there is `x` in X such that `f(x)` is true. If so, returns the first such `x` and its index.}
  idx := 0;
  while true do
    idx +:= 1;
    ok, x := HasNext(X);
    if ok then
      if f(x) then
        return true, x, idx;
      end if;
    else
      return false, _, _;
    end if;
  end while;
end intrinsic;

intrinsic Truncate(X :: PGGIter, n :: RngIntElt) -> PGGIter_Truncate
  {Limits the length of X to n.}
  require n ge 0: "n must be at least 0";
  Y := New(PGGIter_Truncate);
  Y`iter := X;
  Y`limit := n;
  return Y;
end intrinsic;

intrinsic Truncate(X :: PGGIter_List, n :: RngIntElt) -> PGGIter_List
  {"}
  xs := ToList(X);
  if #xs gt n then
    return PGG_ToIter(xs[1..n]);
  else
    return X;
  end if;
end intrinsic;

intrinsic HasNext(X :: PGGIter) -> BoolElt, .
  {True if X has a next item.}
  if X`limit gt 0 then
    ok, x := HasNext(X`iter);
    if ok then
      X`limit -:= 1;
      return true, x;
    end if;
  end if;
  return false, _;
end intrinsic;