/*
 * ProGuard -- shrinking, optimization, obfuscation, and preverification
 *             of Java bytecode.
 *
 * Copyright (c) 2002-2018 GuardSquare NV
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */
package proguard.classfile.visitor;

import proguard.classfile.*;

import java.util.*;
import java.util.concurrent.*;


/**
 * This ClassPoolVisitor will visit all Clazz objects of the class pool
 * in a parallel way. For each thread, a separate ClassVisitor will be
 * created using {@link ClassVisitorFactory#createClassVisitor()}.
 * <p>
 * The number of parallel threads is coupled to the number of available
 * processors:
 * <pre>
 *     parallel_threads = Runtime.getRuntime().availableProcessors() - 1;
 * </pre>
 * <p>
 * It is possible to override the number of threads by setting the
 * environment variable {@code parallel.threads} to an integer > 0.
 *
 * @author Thomas Neidhart
 */
public class ParallelAllClassVisitor
implements ClassPoolVisitor
{
    private static final int THREAD_COUNT;
    static {
        Integer threads = null;
        try {
            String threadCountString = System.getProperty("parallel.threads");
            if (threadCountString != null)
            {
                threads = Integer.parseInt(threadCountString);
            }
        }
        catch (Exception ex) {}

        threads = threads == null ?
            Runtime.getRuntime().availableProcessors() - 1 :
            Math.min(threads.intValue(), Runtime.getRuntime().availableProcessors());

        THREAD_COUNT = threads.intValue();
    }


    /**
     * A factory for ClassVisitor objects.
     */
    public interface ClassVisitorFactory
    {
        /**
         * Creates a ClassVisitor that will be used during
         * parallel visiting of classes in a ClassPool.
         */
        ClassVisitor createClassVisitor();
    }


    private final ClassVisitorFactory  classVisitorFactory;


    /**
     * Create a new ParallelAllClassVisitor that will use the given factory
     * to visit all classes in a ClassPool in a parallel way.
     */
    public ParallelAllClassVisitor(ClassVisitorFactory classVisitorFactory)
    {
        this.classVisitorFactory = classVisitorFactory;
    }


    // Implementations for ClassPoolVisitor.

    public void visitClassPool(ClassPool classPool)
    {
        if (THREAD_COUNT <= 1)
        {
            // Fallback to single thread execution if the thread count
            // was overridden by an environment variable.
            classPool.classesAccept(classVisitorFactory.createClassVisitor());
        }
        else
        {
            ExecutorService executor = Executors.newFixedThreadPool(THREAD_COUNT, new MyThreadFactory());

            MyThreadedClassVisitor classVisitor = new MyThreadedClassVisitor(executor);

            classPool.classesAccept(classVisitor);

            try
            {
                // Shutdown the executor service to release memory.
                executor.shutdown();

                // Rethrow any exception that was thrown in the executor threads.
                classVisitor.awaitTermination();
            }
            catch (InterruptedException e)
            {
                throw new RuntimeException("Parallel execution is taking too long", e);
            }
            catch (ExecutionException e)
            {
                throw new RuntimeException(e.getCause());
            }
        }
    }


    private class MyThreadFactory
    implements ThreadFactory
    {
        private int threadCounter = 0;

        public Thread newThread(Runnable runnable)
        {
            return new MyClassVisitorThread(++threadCounter, runnable);
        }
    }


    private class MyClassVisitorThread
    extends       Thread
    {
        private final ClassVisitor classVisitor = classVisitorFactory.createClassVisitor();

        public MyClassVisitorThread(int counter, Runnable runnable)
        {
            super(runnable, "Parallel Class Visitor " + counter);
        }
    }


    private static class MyThreadedClassVisitor
    implements ClassVisitor
    {
        private final ExecutorService executorService;

        private final List<Future> futures = new ArrayList<Future>();

        public MyThreadedClassVisitor(ExecutorService executorService)
        {
            this.executorService = executorService;
        }

        public void awaitTermination() throws ExecutionException, InterruptedException
        {
            for (Future future : futures)
            {
                future.get();
            }
        }

        // Implementations for ClassVisitor.

        public void visitLibraryClass(LibraryClass libraryClass)
        {
            submitClassToExecutorService(libraryClass);
        }


        public void visitProgramClass(ProgramClass programClass)
        {
            submitClassToExecutorService(programClass);
        }


        private void submitClassToExecutorService(final Clazz clazz)
        {
            futures.add(executorService.submit(new Runnable()
            {
                public void run()
                {
                    MyClassVisitorThread thread = (MyClassVisitorThread)Thread.currentThread();
                    clazz.accept(thread.classVisitor);
                }
            }));
        }
    }
}
