/* Iterator pattern in Chalice. */

class List module Collections {
  var contents: seq<int>;

  method init()
    requires acc(contents);
    ensures valid && size()==0;
  {
    contents := nil<int>;
    fold valid; 
  }

  method add(x: int) 
    requires valid;
    ensures valid && size() == old(size()+1) && get(size()-1) == x; // I don't know why this happens.
    ensures forall i in [0..size()-1] :: get(i) == old(get(i));
  {
    unfold valid;
    contents := contents ++ [x];
    fold valid;
  }

  function get(index: int): int 
    requires rd(valid) && 0<=index && index<size();
  {
    unfolding rd(valid) in contents[index]
  }

  function size(): int 
    requires rd(valid);
    ensures 0<=result;
  {
    unfolding rd(valid) in |contents|
  }

  predicate valid {
    acc(contents)
  }
}

class Iterator module Collections {
  var list: List;
  var index: int;

  method init(l: List)
    requires acc(list) && acc(index);
    requires l!=null;
    requires rd(l.valid);
    ensures valid;
    ensures getList()==l;
  {
    list := l;
    this.index := 0;
    fold valid;
  }

  method next() returns (rt: int) 
    requires valid && hasNext();
    ensures valid;
    ensures getList()==old(getList());
  {
    unfold valid;
    rt := list.get(index);
    index := index + 1;
    fold valid;
  }

  method dispose()
    requires valid;
    ensures rd(old(getList()).valid);
  {
    unfold valid;
  }

  function hasNext(): bool 
    requires valid;
  {
    unfolding valid in index<list.size()
  }

  function getList(): List
    requires valid;
    ensures result!=null;
  {
    unfolding valid in list
  }

  predicate valid
  {
    acc(list) && acc(index) && list!=null && rd(list.valid) && 0<=index && index<=list.size()
  }
}

class Program module Main {
  method main(){
    var tmp: int;
    //create a new list
    var list := new List;
    call list.init();
    call list.add(5);
    call list.add(6);
 
    // create a new iterator
    var iter1 := new Iterator;
    call iter1.init(list);

    // create a second iterator
    var iter2 := new Iterator;
    call iter2.init(list);
    
    // iterate over the list
    while(iter1.hasNext())
      invariant iter1.valid && iter1.getList()==list;
    {
      call tmp := iter1.next();
    }

    // iterate over the list
    while(iter2.hasNext())
      invariant iter2.valid && iter2.getList()==list;
    {
      call tmp := iter2.next();
    }

    // dispose the iterators
    call iter1.dispose();
    call iter2.dispose();

    // full access to the list
    assert list.valid;
    assert list.size()==2;
  }
}