From e74a9487d9fb730b2e2651a1a80b8be9aad91677 Mon Sep 17 00:00:00 2001
From: Daniel Bachhuber <daniel@handbuilt.co>
Date: Wed, 26 Oct 2016 06:54:43 -0700
Subject: [PATCH] Support for profiling `wp_loaded:after`

---
 features/profile-hook.feature | 21 +++++++++++++++++++++
 inc/class-profiler.php        | 13 +++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/features/profile-hook.feature b/features/profile-hook.feature
index 461a711..e337989 100644
--- a/features/profile-hook.feature
+++ b/features/profile-hook.feature
@@ -137,3 +137,24 @@ Feature: Profile a specific hook
       """
       wp-content/mu-plugins/awesome-file.php
       """
+
+  Scenario: Profile the :after hooks
+    Given a WP install
+
+    When I run `wp profile hook wp_loaded:after`
+    Then STDOUT should contain:
+      """
+      do_action()
+      """
+
+    When I run `wp profile hook wp:after`
+    Then STDOUT should contain:
+      """
+      do_action_ref_array()
+      """
+
+    When I run `wp profile hook wp_footer:after`
+    Then STDOUT should contain:
+      """
+      do_action()
+      """
diff --git a/inc/class-profiler.php b/inc/class-profiler.php
index f643a83..78f10bf 100644
--- a/inc/class-profiler.php
+++ b/inc/class-profiler.php
@@ -105,6 +105,10 @@ class Profiler {
 				WP_CLI::add_hook( 'after_wp_config_load', array( $this, 'wp_tick_profile_begin' ) );
 			}
 			WP_CLI::add_wp_hook( $end_hook, array( $this, 'wp_tick_profile_end' ), -9999 );
+		} else if ( 'hook' === $this->type
+			&& ':after' === substr( $this->focus, -6, 6 ) ) {
+			$start_hook = substr( $this->focus, 0, -6 );
+			WP_CLI::add_wp_hook( $start_hook, array( $this, 'wp_tick_profile_begin' ), 9999 );
 		} else {
 			WP_CLI::add_wp_hook( 'all', array( $this, 'wp_hook_begin' ) );
 		}
@@ -387,6 +391,9 @@ class Profiler {
 			$this->loggers[ $this->running_hook ]->stop();
 			$this->running_hook = null;
 		}
+		if ( 'hook' === $this->type && 'wp_loaded:after' === $this->focus ) {
+			$this->wp_tick_profile_end();
+		}
 		if ( 'stage' === $this->type && ! $this->focus ) {
 			$logger->stop();
 			$this->loggers[] = $logger;
@@ -406,6 +413,9 @@ class Profiler {
 			$this->loggers[ $this->running_hook ]->stop();
 			$this->running_hook = null;
 		}
+		if ( 'hook' === $this->type && 'wp:after' === $this->focus ) {
+			$this->wp_tick_profile_end();
+		}
 		if ( 'stage' === $this->type && ! $this->focus ) {
 			$logger->stop();
 			$this->loggers[] = $logger;
@@ -434,6 +444,9 @@ class Profiler {
 			$this->loggers[ $this->running_hook ]->stop();
 			$this->running_hook = null;
 		}
+		if ( 'hook' === $this->type && 'wp_footer:after' === $this->focus ) {
+			$this->wp_tick_profile_end();
+		}
 		if ( 'stage' === $this->type && ! $this->focus ) {
 			$logger->stop();
 			$this->loggers[] = $logger;