From 7f809423d46bc81d3758555cc793aa696a38ef7b Mon Sep 17 00:00:00 2001
From: Stephen Cook <stephen@stephencookdev.co.uk>
Date: Sat, 17 Feb 2018 00:10:58 +0000
Subject: [PATCH] Fix wrapping around plugin callbacks

---
 WrappedPlugin/index.js | 27 +++++++++++++++++----------
 package.json           |  2 +-
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/WrappedPlugin/index.js b/WrappedPlugin/index.js
index 46f52bd..0feeab2 100644
--- a/WrappedPlugin/index.js
+++ b/WrappedPlugin/index.js
@@ -7,8 +7,11 @@ const genPluginMethod = (orig, pluginName, smp, type) =>
       const timeEventName = pluginName + "/" + type + "/" + method;
       // we don't know if there's going to be a callback applied to a particular
       // call, so we just set it multiple times, letting each one override the last
-      const addEndEvent = () =>
+      let endCallCount = 0;
+      const addEndEvent = () => {
+        endCallCount++;
         smp.addTimeEvent("plugins", timeEventName, "end", { id });
+      };
 
       smp.addTimeEvent("plugins", timeEventName, "start", {
         id,
@@ -18,7 +21,11 @@ const genPluginMethod = (orig, pluginName, smp, type) =>
         this,
         args.map(a => wrap(a, pluginName, smp, addEndEvent))
       );
-      addEndEvent();
+
+      // If the end event was invoked as a callback immediately, we can
+      // don't want to add another end event here (and it can actually cause
+      // errors, if webpack has finished compilation entirely)
+      if (!endCallCount) addEndEvent();
 
       return ret;
     };
@@ -45,19 +52,19 @@ const wrap = (orig, pluginName, smp, addEndEvent) => {
     return construcNamesToWrap.includes(origConstrucName);
   };
   const shouldWrap = getShouldWrap(orig);
-  const shouldSoftWrap = Object.values(orig).some(getShouldWrap);
+  const shouldSoftWrap = Object.keys(orig)
+    .map(k => orig[k])
+    .some(getShouldWrap);
 
   if (!shouldWrap && !shouldSoftWrap) {
-    const vanillaFunc =
-      typeof orig === "function" &&
-      orig &&
-      orig.prototype &&
-      orig.prototype.constructor !== orig;
+    const vanillaFunc = orig.name === "next";
     return vanillaFunc
       ? function() {
-          const ret = orig();
+          // do this before calling the callback, since the callback can start
+          // the next plugin step
           addEndEvent();
-          return ret;
+
+          return orig.apply(this, arguments);
         }
       : orig;
   }
diff --git a/package.json b/package.json
index 440a695..ba73a7b 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "speed-measure-webpack-plugin",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "description": "Measure + analyse the speed of your webpack loaders and plugins",
   "main": "index.js",
   "repository": {