import edu.rice.cs.mint.runtime.Code;
import edu.rice.cs.mint.runtime.SafeCode;
import edu.rice.cs.mint.util.Benchmark;

public class unroll2{
    
    /*
     public static void main(String[] args)
     {
     final LoopIteration li = new PrintIteration();
     Code<? extends FibFun> CodeFib6 = <| new FibFun() {
     public void Fib(){
     Integer i;
     Integer init = 0;
     `(unroll( <| i |>, <| init |>, 3, 5, li));
     System.out.println(i);
     }
     } |>;
     FibFun sfib6 = CodeFib6.run();
     sfib6.Fib();
     }
     */
    
    public static void main(String[] args){
        final int start = 0;
        final int stop = 20;
        final int step = 2;
        
        System.out.println("unroll(" + start + "," + stop + "," + step + ")");
        Benchmark.TimedTask[] results =
            Benchmark.benchmark(Benchmark.stagingTasks(new Benchmark.Task() {
            public void run() {
                //
                // call unstaged code here
                //
                final IntCell cell = new IntCell();
                cell.value = 0;
                rolled(start, stop, step, new IncrLoopIteration (cell));
                int j = cell.value;
                //
                //
                //
            }
        }, new Benchmark.Thunk<Code<? extends Benchmark.Task>>() {
            public Code<? extends Benchmark.Task> value() {
                return <| new Benchmark.Task() {
                    public void run() {
                        //
                        // splice in staged code here
                        //
                        final IntCell cell = new IntCell();
                        cell.value = 0;
                        `(unrolled(start, stop, step,
                                   new sIncrLoopIteration (<| cell |>)));
                        int j = cell.value;
                        //
                        //
                        //
                    }
                } |>;
            }
        }));
        
//        Benchmark.print(0, results); // relative to unstaged
//        System.out.println();
        Benchmark.print(results.length-1, results); // relative to staged
        
        
    }
    
    
    public static separable Code<Void> unrolled(int start, int stop, int step,
                                                sLoopIteration F){
        Code<Void> C = <| { } |>;
        for(int x = start; x < stop; x += step){
            C = <| { `C; `( F.iteration(x)); } |>;
        }
        return C;
    }
    
    public static void rolled(int start, int stop, int step, LoopIteration F){
        for(int x = start; x < stop; x += step){
            F.iteration (x);
        }
    }
    
    // FIXME: fix unroll_part work like above
    /*
     public static separable Code<Void> unroll_part(Code<IntCell> i,
     Code<Integer> init,
     int incr,
     int iterations,
     sLoopIteration F,
     int blockSize){
     final int loops = iterations/blockSize;
     final int leftover = iterations % blockSize;
     final int fBlockSize = blockSize;
     final int fIncr = incr;
     if(loops < 2)
     return unroll(i, init, incr, iterations, F);
     else{
     final Code<Integer> ii = <| `(i).value |>;
     return
     <| { for (`(i).value = `init;
     `(i).value < (`init + new Integer(loops * fBlockSize * fIncr));){
     `(unroll(i, ii, fIncr, fBlockSize, F));
     }
     `(unroll(i, ii, fIncr, leftover, F)); } |>;
     }
     }
     */
    
    
    public static class IntCell{
        public Integer value = 0;
    }
    
    static interface LoopIteration{
        public void iteration (int i);
    }
    
    static class IncrLoopIteration implements LoopIteration {
        IntCell cell;
        
        IncrLoopIteration (IntCell cell) {
            this.cell = cell;
        }
        
        public void iteration (int i) {
            cell.value += i;
        }
    }
    
    static class PrintIteration implements LoopIteration{
        public void iteration (int i){
            new Integer(i * i);
        }
    }
    
    static interface sLoopIteration{
        public separable Code<Void> iteration (int i);
    }
    
    static class sIncrLoopIteration implements sLoopIteration {
        Code<IntCell> cell;
        
        sIncrLoopIteration (Code<IntCell> cell) {
            this.cell = cell;
        }
        
        public separable Code<Void> iteration (final int i) {
            return <| { (`cell).value += i; } |>;
        }
    }
    
    static class sPrintIteration implements sLoopIteration{
        public separable Code<Void> iteration (final int i){
            return <| {new Integer(i * i); } |>;
        }
    }
    
}
