/*

  * Copyright (c) 2003-2008, by Henrik Arro and Contributors

  *

  * This file is part of JSeq, a tool to automatically create

  * sequence diagrams by tracing program execution.

  *

  * See <http://jseq.sourceforge.net> for more information.

  *

  * JSeq is free software: you can redistribute it and/or modify

  * it under the terms of the GNU Lesser General Public License as

  * published by the Free Software Foundation, either version 3 of

  * the License, or (at your option) any later version.

  *

  * JSeq 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 Lesser General Public License for more details.

  *

  * You should have received a copy of the GNU Lesser General Public License

  * along with JSeq. If not, see <http://www.gnu.org/licenses/>.

  */

 package th.co.edge.jseq;

 import java.util.HashMap;

 import java.util.List;

 import java.util.Map;

 import com.sun.jdi.AbsentInformationException;

 import com.sun.jdi.IncompatibleThreadStateException;

 import com.sun.jdi.Location;

 import com.sun.jdi.Method;

 import com.sun.jdi.ReferenceType;

 import com.sun.jdi.ThreadReference;

 import com.sun.jdi.VMDisconnectedException;

 import com.sun.jdi.VirtualMachine;

 import com.sun.jdi.event.BreakpointEvent;

 import com.sun.jdi.event.ClassPrepareEvent;

 import com.sun.jdi.event.Event;

 import com.sun.jdi.event.EventIterator;

 import com.sun.jdi.event.EventQueue;

 import com.sun.jdi.event.EventSet;

 import com.sun.jdi.event.ExceptionEvent;

 import com.sun.jdi.event.LocatableEvent;

 import com.sun.jdi.event.MethodEntryEvent;

 import com.sun.jdi.event.MethodExitEvent;

 import com.sun.jdi.event.StepEvent;

 import com.sun.jdi.event.ThreadDeathEvent;

 import com.sun.jdi.event.VMDeathEvent;

 import com.sun.jdi.event.VMDisconnectEvent;

 import com.sun.jdi.event.VMStartEvent;

 import com.sun.jdi.request.BreakpointRequest;

 import com.sun.jdi.request.ClassPrepareRequest;

 import com.sun.jdi.request.EventRequest;

 import com.sun.jdi.request.EventRequestManager;

 import com.sun.jdi.request.ExceptionRequest;

 import com.sun.jdi.request.MethodEntryRequest;

 import com.sun.jdi.request.MethodExitRequest;

 import com.sun.jdi.request.StepRequest;

 import com.sun.jdi.request.ThreadDeathRequest;

 /**

  * An <code>EventThred</code> traces the execution of a Java process, keeping

  * track of all method calls as <code>Activation</code>s in an

  * <code>ActivationList</code> containing all root activations, that is, all

  * method calls that occur "magically" from outside the traced methods, for

  * example from the main method. Each Java thread is traced as a separate root

  * activation.

  */

 public class EventThread extends Thread {

     private final VirtualMachine vm;

     private final ActivationList rootActivations;

     private final List<String> includes;

     private final List<String> excludes;

     private final List<String> boundaryMethods;

     private final boolean trace;

     private String startClassName;

     private String startMethodName;

     private boolean connected = true;

     private boolean vmDied = true;

     private Map<ThreadReference, ThreadTrace> traceMap =

             new HashMap<ThreadReference, ThreadTrace>();

     /**

      * Creates a new <code>EventThread</code> that traces a given JDI

      * <code>VirtualMachine</code>, filling in a list of root activations.

      *

      * <p>

      * A list of included and excluded method names determine which methods to

      * trace. If the list of included method names is empty, all methods except

      * the excluded ones are traced, and correspondingly if the list of excluded

      * method name is empty. If both lists are non-empty, only the methods in

      * the included list are traced, except the ones in the excluded list.

      * Method names in the included and excluded lists can either be a fully

      * qualified method name, or a simple wild-card expression starting or

      * ending with "*".

      *

      * <p>

      * For example, if the included list contains "foo.Bar.*", and the excluded

      * list is empty, all methods in the class <code>foo.Bar</code> will be

      * traced. If the included list is empty and the excluded list contains

      * "java.*", everything but methods in the <code>java</code> package will

      * be traced. Finally, if the included list contains "foo.Bar.*" and the

      * excluded list contains "foo.Bar.baz", all methods in the class

      * <code>foo.Bar</code> will be traced, except for the method

      * <code>baz</code>.

      *

      * <p>

      * It is also possible to specify a list of boundary methods. A <it>bounday

      * method</it> is a method where tracing stops at entry to the method and

      * is resumed when the method exits. This list should contain fully

      * qualified method names, no wild-cards are accepted.

      *

      * <p>

      * For example, if the list of boundary methods contains "foo.Bar.baz",

      * tracing will include the call to <code>baz</code>, but no methods

      * called by <code>baz</code>. Normal tracing is resumed after

      * <code>baz</code> returns.

      *

      * @param vm

      *            the JDI <code>VirtualMachine</code>, a representation of

      *            the Java process to trace

      * @param rootActivations

      *            the list of root activations that will be filled in during the

      *            program trace

      * @param includes

      *            a list of method name pattern that will be included in the

      *            program trace, where name patterns may start or end with the

      *            wild-card "*"

      * @param excludes

      *            a list of method name patterns that will be excluded in the

      *            program trace, where name patterns may start or end with the

      *            wild-card "*"

      * @param boundaryMethods

      *            a list of boundary methods, methods where tracing will stop

      *            during the execution of the method

      * @param trace

      *            if <code>true</code> the program trace (method entry and

      *            exit, for example) will be echoed to <code>System.out</code>

      */

     public EventThread(VirtualMachine vm, ActivationList rootActivations,

             List<String> includes, List<String> excludes,

             List<String> boundaryMethods, boolean trace) {

         super("event-handler");

         this.vm = vm;

         this.rootActivations = rootActivations;

         this.includes = includes;

         this.excludes = excludes;

         this.trace = trace;

         this.boundaryMethods = boundaryMethods;

     }

     /**

      * Handles all events generated by JDI, running for as long as the Java

      * process is active, or until an attached process disconnects the event

      * thread.

      */

     @Override

     public void run() {

         EventQueue queue = vm.eventQueue();

         while (connected) {

             try {

                 EventSet eventSet = queue.remove();

                 EventIterator it = eventSet.eventIterator();

                 while (it.hasNext()) {

                     Event event = it.nextEvent();

                     handleEvent(event);

                 }

                 eventSet.resume();

             } catch (InterruptedException exc) {

                 System.err.println(exc);

             } catch (VMDisconnectedException discExc) {

                 handleDisconnectedException();

                 break;

             } catch (NoSuchMethodException e) {

                 e.printStackTrace();

             }

         }

     }

     /**

      * Creates all JDI event requests, that is, defines the set of events that

      * will be handled by this <code>EventThread</code>.

      *

      * @param startMethod

      *            the fully qualified name of the method where tracing should

      *            start, or <code>null</code> if tracing should start

      *            immediately

      */

     public void setEventRequests(String startMethod) {

         EventRequestManager mgr = vm.eventRequestManager();

         if (startMethod != null) {

             setStartMethodBreakpoints(startMethod);

         } else {

             // I really, really, want to refactor this clumsy implementation by

             // calling a helper method, but this is difficult since the

             // different request types do not implement a common interface

             // defining the addClassFilter and addClassExclusionFilter methods.

             if (includes.isEmpty()) {

                 MethodEntryRequest menr = mgr.createMethodEntryRequest();

                 for (String excludePattern : excludes) {

                     menr.addClassExclusionFilter(excludePattern);

                 }

                 enableRequest(menr);

             } else {

                 for (String includePattern : includes) {

                     MethodEntryRequest menr = mgr.createMethodEntryRequest();

                     menr.addClassFilter(includePattern);

                     for (String excludePattern : excludes) {

                         menr.addClassExclusionFilter(excludePattern);

                     }

                     enableRequest(menr);

                 }

             }

             if (includes.isEmpty()) {

                 MethodExitRequest mexr = mgr.createMethodExitRequest();

                 for (String excludePattern : excludes) {

                     mexr.addClassExclusionFilter(excludePattern);

                 }

                 enableRequest(mexr);

             } else {

                 for (String includePattern : includes) {

                     MethodExitRequest mexr = mgr.createMethodExitRequest();

                     mexr.addClassFilter(includePattern);

                     for (String excludePattern : excludes) {

                         mexr.addClassExclusionFilter(excludePattern);

                     }

                     enableRequest(mexr);

                 }

             }

             if (includes.isEmpty()) {

                 ExceptionRequest exr =

                         mgr.createExceptionRequest(null, true, true);

                 for (String excludePattern : excludes) {

                     exr.addClassExclusionFilter(excludePattern);

                 }

                 enableRequest(exr);

             } else {

                 for (String includePattern : includes) {

                     ExceptionRequest exr =

                             mgr.createExceptionRequest(null, true, true);

                     exr.addClassFilter(includePattern);

                     for (String excludePattern : excludes) {

                         exr.addClassExclusionFilter(excludePattern);

                     }

                     enableRequest(exr);

                 }

             }

             ThreadDeathRequest tdr = mgr.createThreadDeathRequest();

             enableRequest(tdr);

         }

     }

     private void enableRequest(EventRequest request) {

         request.setSuspendPolicy(EventRequest.SUSPEND_ALL);

         request.enable();

     }

     private void setStartMethodBreakpoints(String qualifiedMethodName) {

         int lastDot = qualifiedMethodName.lastIndexOf('.');

         if (lastDot < 0) {

             String message =

                     "Not a qualified method name: \"" + qualifiedMethodName +

                             "\"";

             throw new IllegalArgumentException(message);

         }

         startClassName = qualifiedMethodName.substring(0, lastDot);

         startMethodName = qualifiedMethodName.substring(lastDot + 1);

         List<ReferenceType> startClasses = vm.classesByName(startClassName);

         boolean methodFound = false;

         for (ReferenceType refType : startClasses) {

             List<Method> methods = refType.methodsByName(startMethodName);

             if (!methods.isEmpty()) {

                 setMethodBreakpoints(methods);

                 methodFound = true;

             }

         }

         if (!methodFound) {

             System.err.println("Start method not found: " +

                     qualifiedMethodName +

                     ". Waiting to see if the class will be loaded.");

             EventRequestManager mgr = vm.eventRequestManager();

             ClassPrepareRequest cpr = mgr.createClassPrepareRequest();

             enableRequest(cpr);

         }

     }

     private void setMethodBreakpoints(List<Method> methods) {

         EventRequestManager mgr = vm.eventRequestManager();

         for (Method method : methods) {

             try {

                 Location location = method.allLineLocations().get(0);

                 BreakpointRequest bpr = mgr.createBreakpointRequest(location);

                 enableRequest(bpr);

             } catch (AbsentInformationException e) {

                 e.printStackTrace();

             }

         }

     }

     private void deleteAllEventRequests() {

         EventRequestManager mgr = vm.eventRequestManager();

         mgr.deleteEventRequests(mgr.methodEntryRequests());

         mgr.deleteEventRequests(mgr.methodExitRequests());

         mgr.deleteEventRequests(mgr.exceptionRequests());

     }

     private class ThreadTrace {

         private final ThreadReference thread;

         private Activation currentActivation = null;

         private String currentBoundaryMethod;

         private boolean stopWhenActivationDone = false;

         public ThreadTrace(ThreadReference thread) {

             this.thread = thread;

             trace("====== " + thread.name() + " ======");

         }

         private void classPrepareEvent(ClassPrepareEvent event)

                 throws NoSuchMethodException {

             if (event.referenceType().name().equals(startClassName)) {

                 List<Method> methods =

                         event.referenceType().methodsByName(startMethodName);

                 if (methods.isEmpty()) {

                     throw new NoSuchMethodException(startClassName + "." +

                             startMethodName);

                 }

                 setMethodBreakpoints(methods);

             }

         }

         private void methodEntryEvent(MethodEntryEvent event) {

             String className = getClassName(event);

             Method method = event.method();

             List<String> argumentTypeNames = event.method().argumentTypeNames();

             if (isBoundaryMethod(event)) {

                 deleteAllEventRequests();

                 currentBoundaryMethod = className + "." + method.name();

                 EventRequestManager mgr = vm.eventRequestManager();

                 MethodExitRequest methodExitRequest =

                         mgr.createMethodExitRequest();

                 enableRequest(methodExitRequest);

             }

             methodEntry(className, method, argumentTypeNames);

         }

         private boolean isBoundaryMethod(MethodEntryEvent event) {

             String className = getClassName(event);

             String methodName = event.method().name();

             String qualifiedMethodName = className + "." + methodName;

             for (String boundaryMethod : boundaryMethods) {

                 if (qualifiedMethodName.equals(boundaryMethod)) {

                     return true;

                 }

             }

             return false;

         }

         private void methodEntry(String className, Method method,

                 List<String> argumentTypeNames) {

             trace(className + "." + method.name() + " " + argumentTypeNames);

             int frameCount = -1;

             try {

                 frameCount = thread.frameCount();

             } catch (IncompatibleThreadStateException e) {

                 e.printStackTrace();

                 // Ignore -- I THINK this exception cannot occur here...

             }

             Activation activation =

                     new Activation(currentActivation, className, method,

                             frameCount);

             if (currentActivation == null) {

                 rootActivations.add(activation);

             }

             currentActivation = activation;

         }

         private void methodExitEvent(MethodExitEvent event) {

             if (currentActivation == null) {

                 return;

             }

             String className = getClassName(event);

             String methodName = event.method().name();

             String qualifiedMethodName = className + "." + methodName;

             if (className.equals(currentActivation.getClassName()) &&

                     methodName.equals(currentActivation.getMethod().name())) {

                 currentActivation = currentActivation.getParent();

             }

             if (currentActivation == null && stopWhenActivationDone) {

                 deleteAllEventRequests();

             }

             if (currentBoundaryMethod != null &&

                     currentBoundaryMethod.equals(qualifiedMethodName)) {

                 setEventRequests(null);

             }

         }

         private void exceptionEvent(ExceptionEvent event) {

             EventRequestManager mgr = vm.eventRequestManager();

             StepRequest request =

                     mgr.createStepRequest(thread, StepRequest.STEP_MIN,

                             StepRequest.STEP_INTO);

             request.addCountFilter(1);

             enableRequest(request);

         }

         private void stepEvent(StepEvent event) {

             EventRequestManager mgr = vm.eventRequestManager();

             mgr.deleteEventRequest(event.request());

             // Find the activation that catches the exception.

             int frameCount = -1;

             try {

                 frameCount = thread.frameCount();

             } catch (IncompatibleThreadStateException e) {

                 System.err.println(e);

             }

             while (currentActivation.getParent() != null) {

                 if (currentActivation.getFrameCount() == frameCount) {

                     break;

                 }

                 currentActivation = currentActivation.getParent();

             }

         }

         private void breakpointEvent(BreakpointEvent event) {

             String methodName = event.location().method().name();

             if (!methodName.equals(startMethodName)) {

                 String message =

                         "Unexpected breakpoint. Should be in method " +

                                 startMethodName + ", but was in method " +

                                 methodName + ". Event=" + event;

                 throw new IllegalArgumentException(message);

             }

             methodEntry(startClassName, event.location().method(), event

                     .location().method().argumentTypeNames());

             EventRequestManager mgr = vm.eventRequestManager();

             mgr.deleteEventRequest(event.request());

             setEventRequests(null);

             stopWhenActivationDone = true;

         }

         private void threadDeathEvent(ThreadDeathEvent event) {

             trace("====== " + thread.name() + " end ======");

         }

         private String getClassName(LocatableEvent event) {

             String className = getDeclaringType(event).name();

             try {

                 className = getInstanceType(event).name();

             } catch (IncompatibleThreadStateException itse) {

                 System.err.println("cannot identify instance type for " +

                         "method call: " + event.location().method().name() +

                         "; using declaring type: " + className);

             }

             return className;

         }

         private ReferenceType getDeclaringType(LocatableEvent event) {

             return event.location().method().declaringType();

         }

         private ReferenceType getInstanceType(LocatableEvent event)

                 throws IncompatibleThreadStateException {

             try {

                 return event.thread().frame(0).thisObject().referenceType();

             } catch (Exception e) {

                 throw new IncompatibleThreadStateException(e.getMessage());

             }

         }

     }

     private ThreadTrace threadTrace(ThreadReference thread) {

         ThreadTrace threadTrace = traceMap.get(thread);

         if (threadTrace == null) {

             threadTrace = new ThreadTrace(thread);

             traceMap.put(thread, threadTrace);

         }

         return threadTrace;

     }

     private void handleEvent(Event event) throws NoSuchMethodException {

         if (event instanceof ClassPrepareEvent) {

             classPrepareEvent((ClassPrepareEvent) event);

         } else if (event instanceof MethodEntryEvent) {

             methodEntryEvent((MethodEntryEvent) event);

         } else if (event instanceof MethodExitEvent) {

             methodExitEvent((MethodExitEvent) event);

         } else if (event instanceof ExceptionEvent) {

             exceptionEvent((ExceptionEvent) event);

         } else if (event instanceof StepEvent) {

             stepEvent((StepEvent) event);

         } else if (event instanceof BreakpointEvent) {

             breakpointEvent((BreakpointEvent) event);

         } else if (event instanceof ThreadDeathEvent) {

             threadDeathEvent((ThreadDeathEvent) event);

         } else if (event instanceof VMStartEvent) {

             vmStartEvent((VMStartEvent) event);

         } else if (event instanceof VMDeathEvent) {

             vmDeathEvent((VMDeathEvent) event);

         } else if (event instanceof VMDisconnectEvent) {

             vmDisconnectEvent((VMDisconnectEvent) event);

         } else {

             throw new Error("Internal error: Unexpected event type");

         }

     }

     private synchronized void handleDisconnectedException() {

         EventQueue queue = vm.eventQueue();

         while (connected) {

             try {

                 EventSet eventSet = queue.remove();

                 EventIterator iter = eventSet.eventIterator();

                 while (iter.hasNext()) {

                     Event event = iter.nextEvent();

                     if (event instanceof VMDeathEvent) {

                         vmDeathEvent((VMDeathEvent) event);

                     } else if (event instanceof VMDisconnectEvent) {

                         vmDisconnectEvent((VMDisconnectEvent) event);

                     }

                 }

                 eventSet.resume();

             } catch (InterruptedException e) {

                 System.err.println(e);

             }

         }

     }

     private void vmStartEvent(VMStartEvent event) {

         trace("-- VM Started --");

     }

     private void classPrepareEvent(ClassPrepareEvent event)

             throws NoSuchMethodException {

         threadTrace(event.thread()).classPrepareEvent(event);

     }

     private void methodEntryEvent(MethodEntryEvent event) {

         threadTrace(event.thread()).methodEntryEvent(event);

     }

     private void methodExitEvent(MethodExitEvent event) {

         threadTrace(event.thread()).methodExitEvent(event);

     }

     private void stepEvent(StepEvent event) {

         threadTrace(event.thread()).stepEvent(event);

     }

     private void breakpointEvent(BreakpointEvent event) {

         threadTrace(event.thread()).breakpointEvent(event);

     }

     private void threadDeathEvent(ThreadDeathEvent event) {

         ThreadTrace threadTrace = traceMap.get(event.thread());

         if (threadTrace != null) {

             threadTrace.threadDeathEvent(event);

         }

     }

     private void exceptionEvent(ExceptionEvent event) {

         ThreadTrace threadTrace = traceMap.get(event.thread());

         if (threadTrace != null) {

             threadTrace.exceptionEvent(event);

         }

     }

     private void vmDeathEvent(VMDeathEvent event) {

         vmDied = true;

         trace("-- The application exited --");

     }

     private void vmDisconnectEvent(VMDisconnectEvent event) {

         connected = false;

         if (!vmDied) {

             trace("-- The application has been disconnected --");

         }

     }

     private void trace(String s) {

         if (trace) {

             System.err.println(s);

         }

     }

 }