/*
 * Copyright (c) 2013, Regents of the University of California
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package edu.uci.python.nodes.generator;

import com.oracle.truffle.api.*;
import com.oracle.truffle.api.frame.*;
import com.oracle.truffle.api.nodes.*;

import edu.uci.python.nodes.*;
import edu.uci.python.nodes.control.*;
import edu.uci.python.runtime.datatype.*;
import edu.uci.python.runtime.exception.*;
import edu.uci.python.runtime.function.*;
import edu.uci.python.runtime.iterator.*;

public class GeneratorForNode extends LoopNode {

    @Child protected WriteGeneratorFrameVariableNode target;
    @Child protected GetIteratorNode getIterator;

    private final int iteratorSlot;
    private int count;

    public GeneratorForNode(WriteGeneratorFrameVariableNode target, GetIteratorNode getIterator, PNode body, int iteratorSlot) {
        super(body);
        this.target = target;
        this.getIterator = getIterator;
        this.iteratorSlot = iteratorSlot;
    }

    public int getIteratorSlot() {
        return iteratorSlot;
    }

    protected PIterator getIterator(VirtualFrame frame) {
        return PArguments.getGeneratorArguments(frame).getIteratorAt(iteratorSlot);
    }

    protected void setIterator(VirtualFrame frame, PIterator value) {
        PArguments.getGeneratorArguments(frame).setIteratorAt(iteratorSlot, value);
    }

    protected void reset(VirtualFrame frame) {
        setIterator(frame, null);
        count = 0;
    }

    protected void incrementCounter() {
        if (CompilerDirectives.inInterpreter()) {
            count++;
        }
    }

    @Override
    public Object execute(VirtualFrame frame) {
        try {
            executeIterator(frame);
        } catch (StopIterationException e) {
            reset(frame);
            return PNone.NONE;
        }

        try {
            while (true) {
                body.executeVoid(frame);
                target.executeWith(frame, getIterator(frame).__next__());
                incrementCounter();
            }
        } catch (StopIterationException e) {
            if (CompilerDirectives.inInterpreter()) {
                reportLoopCount(count);
            }
        }

        reset(frame);
        return PNone.NONE;
    }

    protected void executeIterator(VirtualFrame frame) throws StopIterationException {
        if (getIterator(frame) != null) {
            return;
        }

        try {
            setIterator(frame, getIterator.executePIterator(frame));
        } catch (UnexpectedResultException e) {
            throw new RuntimeException();
        }

        target.executeWith(frame, getIterator(frame).__next__());
        incrementCounter();
    }

}

@@ -1788,16 +1788,22 @@
 %0Apublic 
+final 
 class Ge
@@ -2355,23 +2355,21 @@
 %0A%0A    pr
-otec
+iva
 te
-d
  PIterat
@@ -2498,31 +2498,29 @@
    %7D%0A%0A    pr
-otec
+iva
 te
-d
  void setIte
@@ -2665,26 +2665,29 @@
   pr
-otected void reset
+ivate Object doReturn
 (Vir
@@ -2758,16 +2758,43 @@
 nt = 0;%0A
+        return PNone.NONE;%0A
     %7D%0A%0A 
@@ -2794,31 +2794,29 @@
    %7D%0A%0A    pr
-otec
+iva
 te
-d
  void increm
@@ -2914,24 +2914,162 @@
    %7D%0A    %7D%0A%0A
+    private void reportCounter() %7B%0A        if (CompilerDirectives.inInterpreter()) %7B%0A            reportLoopCount(count);%0A        %7D%0A    %7D%0A%0A
     @Overrid
@@ -3231,49 +3231,28 @@
   re
-set(frame);%0A            return PNone.NONE
+turn doReturn(frame)
 ;%0A  
@@ -3515,66 +3515,8 @@
 ) %7B%0A
-            if (CompilerDirectives.inInterpreter()) %7B%0A    
     
@@ -3525,39 +3525,32 @@
       report
-Loop
 Count
-(count
+er(
 );%0A         
@@ -3552,72 +3552,41 @@
     
-    %7D%0A        %7D%0A%0A        reset(frame);%0A        return PNone.NONE
+%7D%0A%0A        return doReturn(frame)
 ;%0A  
@@ -3596,23 +3596,21 @@
 %0A%0A    pr
-otec
+iva
 te
-d
  void ex

Clean up in GeneratorForNode