diff --git a/Android.bp b/Android.bp
index 78939cf..724d479 100644
--- a/Android.bp
+++ b/Android.bp
@@ -12384,6 +12384,7 @@
         "src/trace_processor/perfetto_sql/stdlib/sched/runnable.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/states.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql",
+        "src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span_with_slice.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_level_parallelism.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_state_flattened.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/time_in_state.sql",
diff --git a/BUILD b/BUILD
index f355d4e..b8f0a29 100644
--- a/BUILD
+++ b/BUILD
@@ -2531,6 +2531,7 @@
         "src/trace_processor/perfetto_sql/stdlib/sched/runnable.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/states.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql",
+        "src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span_with_slice.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_level_parallelism.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/thread_state_flattened.sql",
         "src/trace_processor/perfetto_sql/stdlib/sched/time_in_state.sql",
diff --git a/src/trace_processor/perfetto_sql/stdlib/sched/BUILD.gn b/src/trace_processor/perfetto_sql/stdlib/sched/BUILD.gn
index d9b8f42..9168926 100644
--- a/src/trace_processor/perfetto_sql/stdlib/sched/BUILD.gn
+++ b/src/trace_processor/perfetto_sql/stdlib/sched/BUILD.gn
@@ -20,6 +20,7 @@
     "runnable.sql",
     "states.sql",
     "thread_executing_span.sql",
+    "thread_executing_span_with_slice.sql",
     "thread_level_parallelism.sql",
     "thread_state_flattened.sql",
     "time_in_state.sql",
diff --git a/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql b/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql
index ffafd6d..e3a92e2 100644
--- a/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql
+++ b/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span.sql
@@ -14,7 +14,7 @@
 -- limitations under the License.
 --
 
-INCLUDE PERFETTO MODULE slices.flat_slices;
+INCLUDE PERFETTO MODULE graphs.search;
 
 -- A 'thread_executing_span' is thread_state span starting with a runnable slice
 -- until the next runnable slice that's woken up by a process (as opposed
@@ -33,7 +33,7 @@
 -- so this table might contain wakeups from interrupt context, consequently, the
 -- wakeup graph generated might not be accurate.
 --
-CREATE PERFETTO VIEW _runnable_state
+CREATE PERFETTO TABLE _runnable_state
 AS
 SELECT
   thread_state.id,
@@ -50,7 +50,7 @@
   AND (thread_state.irq_context = 0 OR thread_state.irq_context IS NULL);
 
 -- Similar to |_runnable_state| but finds the first runnable state at thread.
-CREATE PERFETTO VIEW _first_runnable_state
+CREATE PERFETTO TABLE _first_runnable_state
 AS
 WITH
   first_state AS (
@@ -77,7 +77,7 @@
 
 --
 -- Finds all sleep states including interruptible (S) and uninterruptible (D).
-CREATE PERFETTO VIEW _sleep_state
+CREATE PERFETTO TABLE _sleep_state
 AS
 SELECT
   thread_state.id,
@@ -92,7 +92,7 @@
 --
 -- Finds the last execution for every thread to end executing_spans without a Sleep.
 --
-CREATE PERFETTO VIEW _thread_end_ts
+CREATE PERFETTO TABLE _thread_end_ts
 AS
 SELECT
   MAX(ts) + dur AS end_ts,
@@ -102,7 +102,7 @@
 GROUP BY utid;
 
 -- Similar to |_sleep_state| but finds the first sleep state in a thread.
-CREATE PERFETTO VIEW _first_sleep_state
+CREATE PERFETTO TABLE _first_sleep_state
 AS
 SELECT
   MIN(s.id) AS id,
@@ -145,43 +145,45 @@
 -- end_ts           = S1_ts.
 CREATE PERFETTO TABLE _wakeup
 AS
+WITH
+  all_wakeups AS (
+    SELECT
+      s.state,
+      s.blocked_function,
+      r.id,
+      r.ts AS ts,
+      r.utid AS utid,
+      r.waker_id,
+      r.waker_utid,
+      s.ts AS prev_end_ts
+    FROM _runnable_state r
+    JOIN _sleep_state s
+      ON s.utid = r.utid AND (s.ts + s.dur = r.ts)
+    UNION ALL
+    SELECT
+      NULL AS state,
+      NULL AS blocked_function,
+      r.id,
+      r.ts,
+      r.utid AS utid,
+      r.waker_id,
+      r.waker_utid,
+      NULL AS prev_end_ts
+    FROM _first_runnable_state r
+    LEFT JOIN _first_sleep_state s
+      ON s.utid = r.utid
+  )
 SELECT
-  s.state,
-  s.blocked_function,
-  r.id,
-  r.ts AS ts,
-  r.utid AS utid,
-  r.waker_id,
-  r.waker_utid,
-  IFNULL(LEAD(s.ts) OVER (PARTITION BY r.utid ORDER BY r.ts), thread_end.end_ts) AS end_ts,
-  s.ts AS prev_end_ts,
-  LAG(r.id) OVER (PARTITION BY r.utid ORDER BY r.ts) AS prev_id
-FROM _runnable_state r
-JOIN _sleep_state s
-  ON s.utid = r.utid AND (s.ts + s.dur = r.ts)
-LEFT JOIN _thread_end_ts thread_end
-  USING (utid)
-UNION ALL
-SELECT
-  NULL AS state,
-  NULL AS blocked_function,
-  r.id,
-  r.ts,
-  r.utid AS utid,
-  r.waker_id,
-  r.waker_utid,
-  IFNULL(s.ts, thread_end.end_ts) AS end_ts,
-  NULL AS prev_end_ts,
-  NULL AS prev_id
-FROM _first_runnable_state r
-LEFT JOIN _first_sleep_state s
-  ON s.utid = r.utid
+  all_wakeups.*,
+  LAG(id) OVER (PARTITION BY utid ORDER BY ts) AS prev_id,
+  IFNULL(LEAD(prev_end_ts) OVER (PARTITION BY utid ORDER BY ts), thread_end.end_ts) AS end_ts
+FROM all_wakeups
 LEFT JOIN _thread_end_ts thread_end
   USING (utid);
 
 -- Mapping from running thread state to runnable
 -- TODO(zezeozue): Switch to use `sched_previous_runnable_on_thread`.
-CREATE PERFETTO TABLE _waker_map
+CREATE PERFETTO TABLE _wakeup_map
 AS
 WITH x AS (
 SELECT id, waker_id, utid, state FROM thread_state WHERE state = 'Running' AND dur != -1
@@ -193,12 +195,12 @@
     SELECT
       id AS waker_id,
       state,
-      max(id)
+      MAX(id)
         filter(WHERE state = 'R')
           OVER (PARTITION BY utid ORDER BY id) AS id
     FROM x
   )
-SELECT id, waker_id FROM y WHERE state = 'Running' AND id IS NOT NULL ORDER BY waker_id;
+SELECT id, waker_id FROM y WHERE state = 'Running' ORDER BY waker_id;
 
 --
 -- Builds the parent-child chain from all thread_executing_spans. The parent is the waker and
@@ -207,10 +209,10 @@
 -- Note that this doesn't include the roots. We'll compute the roots below.
 -- This two step process improves performance because it's more efficient to scan
 -- parent and find a child between than to scan child and find the parent it lies between.
-CREATE PERFETTO TABLE _wakeup_chain
+CREATE PERFETTO TABLE _wakeup_graph
 AS
 SELECT
-  _waker_map.id AS parent_id,
+  _wakeup_map.id AS waker_id,
   prev_id,
   prev_end_ts,
   _wakeup.id AS id,
@@ -221,178 +223,259 @@
   _wakeup.state,
   _wakeup.blocked_function
 FROM _wakeup
-JOIN _waker_map USING(waker_id);
+JOIN _wakeup_map USING(waker_id)
+ORDER BY id;
 
 -- The inverse of thread_executing_spans. All the sleeping periods between thread_executing_spans.
-CREATE PERFETTO TABLE _sleeping_span
+CREATE PERFETTO TABLE _sleep
 AS
 WITH
   x AS (
     SELECT
       id,
       ts,
-      lag(end_ts) OVER (PARTITION BY utid ORDER BY id) AS prev_end_ts,
-      utid
-    FROM _wakeup_chain
+      prev_end_ts,
+      utid,
+      state,
+      blocked_function
+    FROM _wakeup_graph
   )
 SELECT
-  ts - prev_end_ts AS dur, prev_end_ts AS ts, id AS root_id, utid AS root_utid
+  ts - prev_end_ts AS dur,
+  prev_end_ts AS ts,
+  id AS root_node_id,
+  utid AS critical_path_utid,
+  id AS critical_path_id,
+  ts - prev_end_ts AS critical_path_blocked_dur,
+  state AS critical_path_blocked_state,
+  blocked_function AS critical_path_blocked_function
 FROM x
 WHERE ts IS NOT NULL;
 
---
--- Finds the roots of the |_wakeup_chain|.
-CREATE PERFETTO TABLE _wakeup_root
-AS
-WITH
-  _wakeup_root_id AS (
-    SELECT DISTINCT parent_id AS id FROM _wakeup_chain
-    EXCEPT
-    SELECT DISTINCT id FROM _wakeup_chain
+-- Given a set of critical paths identified by their |root_node_ids|, flattens
+-- the critical path tasks such that there are no overlapping intervals. The end of a
+-- task in the critical path is the start of the following task in the critical path.
+CREATE PERFETTO MACRO _flatten_critical_path_tasks(_critical_path_table TableOrSubquery)
+RETURNS TableOrSubquery
+AS (
+  WITH
+    x AS (
+      SELECT
+        LEAD(ts) OVER (PARTITION BY root_node_id ORDER BY node_id) AS ts,
+        node_id,
+        ts AS node_ts,
+        root_node_id,
+        utid AS node_utid,
+        _wakeup_graph.prev_end_ts
+      FROM $_critical_path_table
+      JOIN _wakeup_graph
+        ON node_id = id
+    )
+  SELECT node_ts AS ts, root_node_id, node_id, ts - node_ts AS dur, node_utid, prev_end_ts FROM x
+);
+
+-- Converts a table with <ts, dur, utid> columns to a unique set of wakeup roots <id> that
+-- completely cover the time intervals.
+CREATE PERFETTO MACRO _intervals_to_roots(source_table TableOrSubQuery)
+RETURNS TableOrSubQuery
+AS (
+  WITH source AS (
+    SELECT * FROM $source_table
+  ), thread_bounds AS (
+    SELECT utid, MIN(ts) AS min_start, MAX(ts) AS max_start FROM _wakeup_graph GROUP BY utid
+  ), start AS (
+    SELECT
+      _wakeup_graph.utid, max(_wakeup_graph.id) AS start_id, source.ts, source.dur
+      FROM _wakeup_graph
+      JOIN thread_bounds
+        USING (utid)
+      JOIN source
+        ON source.utid = _wakeup_graph.utid AND MAX(source.ts, min_start) >= _wakeup_graph.ts
+     GROUP BY source.ts, source.utid
+  ), end AS (
+    SELECT
+      _wakeup_graph.utid, min(_wakeup_graph.id) AS end_id, source.ts, source.dur
+      FROM _wakeup_graph
+      JOIN thread_bounds
+          USING (utid)
+      JOIN source ON source.utid = _wakeup_graph.utid
+          AND MIN((source.ts + source.dur), max_start) <= _wakeup_graph.ts
+     GROUP BY source.ts, source.utid
+  ), bound AS (
+    SELECT start.utid, start.ts, start.dur, start_id, end_id
+      FROM start
+      JOIN end ON start.ts = end.ts AND start.dur = end.dur AND start.utid = end.utid
   )
-SELECT NULL AS parent_id, _wakeup.*
-FROM _wakeup
-JOIN _wakeup_root_id USING(id);
+  SELECT DISTINCT _wakeup_graph.id FROM bound
+  JOIN _wakeup_graph ON _wakeup_graph.id BETWEEN start_id AND end_id
+);
 
---
--- Finds the leafs of the |_wakeup_chain|.
-CREATE PERFETTO TABLE _wakeup_leaf AS
-WITH
-  _wakeup_leaf_id AS (
-    SELECT DISTINCT id AS id FROM _wakeup_chain
-    EXCEPT
-    SELECT DISTINCT parent_id AS id FROM _wakeup_chain
-  )
-SELECT _wakeup_chain.*
-FROM _wakeup_chain
-JOIN _wakeup_leaf_id USING(id);
-
---
--- Merges the roots, leafs and the rest of the chain.
-CREATE TABLE _wakeup_graph
-AS
-SELECT
-  _wakeup_chain.parent_id,
-  _wakeup_chain.id,
-  _wakeup_chain.ts,
-  _wakeup_chain.end_ts - _wakeup_chain.ts AS dur,
-  _wakeup_chain.utid,
-  _wakeup_chain.prev_end_ts,
-  _wakeup_chain.state,
-  _wakeup_chain.blocked_function,
-  0 AS is_root,
-  (_wakeup_leaf.id IS NOT NULL) AS is_leaf
-FROM _wakeup_chain
-LEFT JOIN _wakeup_leaf
-  USING (id)
-UNION ALL
-SELECT
-  _wakeup_root.parent_id,
-  _wakeup_root.id,
-  _wakeup_root.ts,
-  _wakeup_root.end_ts - _wakeup_root.ts AS dur,
-  _wakeup_root.utid,
-  _wakeup_root.prev_end_ts,
-  _wakeup_root.state,
-  _wakeup_root.blocked_function,
-  1 AS is_root,
-  0 AS is_leaf
-FROM _wakeup_root;
-
--- Thread_executing_span graph of all wakeups across all processes.
---
--- @column root_id            Id of thread_executing_span that initiated the wakeup of |id|.
--- @column parent_id          Id of thread_executing_span that directly woke |id|.
--- @column id                 Id of the first (runnable) thread state in thread_executing_span.
--- @column ts                 Timestamp of first thread_state in thread_executing_span.
--- @column dur                Duration of thread_executing_span.
--- @column utid               Utid of thread with thread_state.
--- @column blocked_dur        Duration of blocking thread state before waking up.
--- @column blocked_state      Thread state ('D' or 'S') of blocked thread_state before waking up.
--- @column blocked_function   Kernel blocked_function of thread state before waking up.
--- @column is_root            Whether the thread_executing_span is a root.
--- @column depth              Tree depth of thread executing span from the root.
-CREATE TABLE _thread_executing_span_graph AS
-WITH roots AS (
-SELECT
-  id AS root_id,
-  parent_id,
-  id,
-  ts,
-  end_ts - ts AS dur,
-  utid,
-  ts - prev_end_ts AS blocked_dur,
-  state AS blocked_state,
-  blocked_function AS blocked_function,
-  1 AS is_root,
-  0 AS depth
-FROM _wakeup_root
-), chain AS (
-  SELECT * FROM roots
-  UNION ALL
+-- Flattens overlapping tasks within a critical path and flattens overlapping critical paths.
+CREATE PERFETTO MACRO _flatten_critical_paths(critical_path_table TableOrSubquery, sleeping_table TableOrSubquery)
+RETURNS TableOrSubquery
+AS (
+  WITH
+    span_starts AS (
+      SELECT
+        cr.node_utid AS utid,
+        MAX(cr.ts, sleep.ts) AS ts,
+        sleep.ts + sleep.dur AS sleep_end_ts,
+        cr.ts + cr.dur AS cr_end_ts,
+        cr.node_id AS id,
+        cr.root_node_id AS root_id,
+        cr.prev_end_ts AS prev_end_ts,
+        critical_path_utid,
+        critical_path_id,
+        critical_path_blocked_dur,
+        critical_path_blocked_state,
+        critical_path_blocked_function
+      FROM
+        _flatten_critical_path_tasks!($critical_path_table) cr
+      JOIN $sleeping_table sleep
+        USING (root_node_id)
+    )
   SELECT
-    chain.root_id,
-    graph.parent_id,
-    graph.id,
-    graph.ts,
-    graph.dur,
-    graph.utid,
-    graph.ts - graph.prev_end_ts AS blocked_dur,
-    graph.state AS blocked_state,
-    graph.blocked_function AS blocked_function,
-    0 AS is_root,
-    chain.depth + 1 AS depth
-  FROM _wakeup_graph graph
-  JOIN chain ON chain.id = graph.parent_id
-) SELECT chain.*, thread.upid FROM chain LEFT JOIN thread USING(utid);
-
--- It finds the MAX between the start of the critical span and the start
--- of the blocked region. This ensures that the critical path doesn't overlap
--- the preceding thread_executing_span before the blocked region.
-CREATE PERFETTO FUNCTION _critical_path_start_ts(ts LONG, leaf_blocked_ts LONG)
-RETURNS LONG AS SELECT MAX($ts, IFNULL($leaf_blocked_ts, $ts));
-
--- See |_thread_executing_span_critical_path|
-CREATE PERFETTO TABLE _critical_path
-AS
-WITH chain AS (
-  SELECT
-    parent_id,
+    ts,
+    MIN(cr_end_ts, sleep_end_ts) - ts AS dur,
+    utid,
     id,
+    root_id,
+    prev_end_ts,
+    critical_path_utid,
+    critical_path_id,
+    critical_path_blocked_dur,
+    critical_path_blocked_state,
+    critical_path_blocked_function
+  FROM span_starts
+  WHERE MIN(sleep_end_ts, cr_end_ts) - ts > 0
+);
+
+-- Generates a critical path.
+CREATE PERFETTO MACRO _critical_path(
+        graph_table TableOrSubquery, root_table TableOrSubquery, sleeping_table TableOrSubquery)
+RETURNS TableOrSubquery
+AS (
+  WITH
+    critical_path AS (
+      SELECT * FROM graph_reachable_weight_bounded_dfs !($graph_table, $root_table, 1)
+    )
+  SELECT
     ts,
     dur,
+    root_id,
+    id,
     utid,
-    id AS critical_path_id,
-    ts - blocked_dur AS critical_path_blocked_ts,
-    blocked_dur AS critical_path_blocked_dur,
-    blocked_state AS critical_path_blocked_state,
-    blocked_function AS critical_path_blocked_function,
-    utid AS critical_path_utid,
-    upid AS critical_path_upid
-  FROM _thread_executing_span_graph graph
+    critical_path_utid,
+    critical_path_id,
+    critical_path_blocked_dur,
+    critical_path_blocked_state,
+    critical_path_blocked_function
+  FROM _flatten_critical_paths!(critical_path, $sleeping_table)
   UNION ALL
+  -- Add roots
   SELECT
-    graph.parent_id,
-    graph.id,
-    _critical_path_start_ts(graph.ts, chain.critical_path_blocked_ts) AS ts,
-    MIN(_critical_path_start_ts(graph.ts, chain.critical_path_blocked_ts) + graph.dur, chain.ts)
-      - _critical_path_start_ts(graph.ts, chain.critical_path_blocked_ts) AS dur,
-    graph.utid,
-    chain.critical_path_id,
-    chain.critical_path_blocked_ts,
-    chain.critical_path_blocked_dur,
-    chain.critical_path_blocked_state,
-    chain.critical_path_blocked_function,
-    chain.critical_path_utid,
-    chain.critical_path_upid
-  FROM _thread_executing_span_graph graph
-  JOIN chain ON (chain.parent_id = graph.id AND (chain.ts > chain.critical_path_blocked_ts))
-) SELECT * FROM chain;
+    ts,
+    end_ts - ts AS dur,
+    id AS root_id,
+    id,
+    utid,
+    utid AS critical_path_utid,
+    NULL AS critical_path_id,
+    NULL AS critical_path_blocked_dur,
+    NULL AS critical_path_blocked_state,
+    NULL AS critical_path_blocked_function
+  FROM $root_table
+  ORDER BY root_id
+);
 
--- Thread executing span critical paths for all threads. For each thread, the critical path of
--- every sleeping thread state is computed and unioned with the thread executing spans on that thread.
--- The duration of a thread executing span in the critical path is the range between the start of the
--- thread_executing_span and the start of the next span in the critical path.
+-- Generates the critical path for only the set of roots <id> passed in.
+-- _intervals_to_roots can be used to generate root ids from a given time interval.
+-- This can be used to genrate the critical path over sparse regions of a trace, e.g
+-- binder transactions. It might be more efficient to generate the _critical_path
+-- for the entire trace, see _thread_executing_span_critical_path_all, but for a
+-- per-process susbset of binder txns for instance, this is likely faster.
+CREATE PERFETTO MACRO _critical_path_by_roots(roots_table TableOrSubQuery)
+RETURNS TableOrSubQuery
+AS (
+  WITH roots AS (
+    SELECT * FROM $roots_table
+  ), root_bounds AS (
+    SELECT MIN(id) AS min_root_id, MAX(id) AS max_root_id FROM roots
+  ), wakeup_bounds AS (
+    SELECT COALESCE(_wakeup_graph.prev_id, min_root_id) AS min_wakeup, max_root_id AS max_wakeup
+    FROM root_bounds
+    JOIN _wakeup_graph ON id = min_root_id
+  ) SELECT
+      id,
+      ts,
+      dur,
+      utid,
+      critical_path_id,
+      critical_path_blocked_dur,
+      critical_path_blocked_state,
+      critical_path_blocked_function,
+      critical_path_utid
+      FROM
+        _critical_path
+        !(
+          (
+            SELECT
+              id AS source_node_id,
+              COALESCE(waker_id, id) AS dest_node_id,
+              id - COALESCE(waker_id, id) AS edge_weight
+            FROM _wakeup_graph
+            JOIN wakeup_bounds WHERE id BETWEEN min_wakeup AND max_wakeup
+          ),
+          (
+            SELECT
+              _wakeup_graph.id AS root_node_id,
+              _wakeup_graph.id - COALESCE(prev_id, _wakeup_graph.id) AS root_target_weight,
+              id,
+              ts,
+              end_ts,
+              utid
+            FROM _wakeup_graph
+            JOIN (SELECT * FROM roots) USING (id)
+          ),
+          _sleep));
+
+-- Generates the critical path for only the time intervals for the utids given.
+-- Currently expensive because of naive interval_intersect implementation.
+-- Prefer _critical_paths_by_roots for performance. This is useful for a small
+-- set of intervals, e.g app startups in a trace.
+CREATE PERFETTO MACRO _critical_path_by_intervals(intervals_table TableOrSubQuery)
+RETURNS TableOrSubQuery AS (
+WITH span_starts AS (
+    SELECT
+      id,
+      MAX(span.ts, intervals.ts) AS ts,
+      MIN(span.ts + span.dur, intervals.ts + intervals.dur) AS end_ts,
+      span.utid,
+      critical_path_id,
+      critical_path_blocked_dur,
+      critical_path_blocked_state,
+      critical_path_blocked_function,
+      critical_path_utid
+    FROM _critical_path_by_roots!(_intervals_to_roots!($intervals_table)) span
+    -- TODO(zezeozue): Replace with interval_intersect when partitions are supported
+    JOIN (SELECT * FROM $intervals_table) intervals ON span.critical_path_utid = intervals.utid
+        AND ((span.ts BETWEEN intervals.ts AND intervals.ts + intervals.dur)
+             OR (intervals.ts BETWEEN span.ts AND span.ts + span.dur))
+) SELECT
+      id,
+      ts,
+      end_ts - ts AS dur,
+      utid,
+      critical_path_id,
+      critical_path_blocked_dur,
+      critical_path_blocked_state,
+      critical_path_blocked_function,
+      critical_path_utid
+   FROM span_starts);
+
+-- Generates the critical path for a given utid over the <ts, dur> interval.
+-- The duration of a thread executing span in the critical path is the range between the
+-- start of the thread_executing_span and the start of the next span in the critical path.
 CREATE PERFETTO FUNCTION _thread_executing_span_critical_path(
   -- Utid of the thread to compute the critical path for.
   critical_path_utid INT,
@@ -420,598 +503,41 @@
   -- Thread Utid the critical path was filtered to.
   critical_path_utid INT
 ) AS
-WITH span_starts AS (
-    SELECT
-      id,
-      MAX(ts, $ts) AS ts,
-      MIN(ts + dur, $ts + $dur) AS end_ts,
-      utid,
-      critical_path_id,
-      critical_path_blocked_dur,
-      critical_path_blocked_state,
-      critical_path_blocked_function,
-      critical_path_utid
-    FROM _critical_path span
-    WHERE (($critical_path_utid IS NOT NULL AND span.critical_path_utid = $critical_path_utid) OR ($critical_path_utid IS NULL))
-      AND ((ts BETWEEN $ts AND $ts + $dur) OR ($ts BETWEEN ts AND ts + dur))
-) SELECT
-      id,
-      ts,
-      end_ts - ts AS dur,
-      utid,
-      critical_path_id,
-      critical_path_blocked_dur,
-      critical_path_blocked_state,
-      critical_path_blocked_function,
-      critical_path_utid
-   FROM span_starts;
+SELECT * FROM _critical_path_by_intervals!((SELECT $critical_path_utid AS utid, $ts as ts, $dur AS dur));
 
--- Limited thread_state view that will later be span joined with the |_thread_executing_span_graph|.
-CREATE PERFETTO VIEW _span_thread_state_view
-AS SELECT id AS thread_state_id, ts, dur, utid, state, blocked_function as function, io_wait, cpu FROM thread_state;
-
--- |_thread_executing_span_graph| span joined with thread_state information.
-CREATE VIRTUAL TABLE _span_graph_thread_state_sp
-USING
-  SPAN_JOIN(
-    _thread_executing_span_graph PARTITIONED utid,
-    _span_thread_state_view PARTITIONED utid);
-
--- Limited slice_view that will later be span joined with the |_thread_executing_span_graph|.
-CREATE PERFETTO VIEW _span_slice_view
+-- Generates the critical path for all threads for the entire trace duration.
+-- The duration of a thread executing span in the critical path is the range between the
+-- start of the thread_executing_span and the start of the next span in the critical path.
+CREATE PERFETTO FUNCTION _thread_executing_span_critical_path_all()
+RETURNS
+  TABLE(
+    -- Id of the first (runnable) thread state in thread_executing_span.
+    id INT,
+    -- Timestamp of first thread_state in thread_executing_span.
+    ts LONG,
+    -- Duration of thread_executing_span.
+    dur LONG,
+    -- Utid of thread with thread_state.
+    utid INT,
+    -- Id of thread executing span following the sleeping thread state for which the critical path is computed.
+    critical_path_id INT,
+    -- Critical path duration.
+    critical_path_blocked_dur LONG,
+    -- Sleeping thread state in critical path.
+    critical_path_blocked_state STRING,
+    -- Kernel blocked_function of the critical path.
+    critical_path_blocked_function STRING,
+    -- Thread Utid the critical path was filtered to.
+    critical_path_utid INT)
 AS
 SELECT
-  slice_id,
-  depth AS slice_depth,
-  name AS slice_name,
-  CAST(ts AS INT) AS ts,
-  CAST(dur AS INT) AS dur,
-  utid
-FROM _slice_flattened;
-
--- |_thread_executing_span_graph| span joined with slice information.
-CREATE VIRTUAL TABLE _span_graph_slice_sp
-USING
-  SPAN_JOIN(
-    _thread_executing_span_graph PARTITIONED utid,
-    _span_slice_view PARTITIONED utid);
-
--- Limited |_thread_executing_span_graph| + thread_state view.
-CREATE PERFETTO VIEW _span_graph_thread_state
-AS
-SELECT ts, dur, id, thread_state_id, state, function, io_wait, cpu
-FROM _span_graph_thread_state_sp;
-
--- Limited |_thread_executing_span_graph| + slice view.
-CREATE PERFETTO VIEW _span_graph_slice
-AS
-SELECT ts, dur, id, slice_id, slice_depth, slice_name
-FROM _span_graph_slice_sp;
-
--- |_thread_executing_span_graph| + thread_state view joined with critical_path information.
-CREATE PERFETTO TABLE _critical_path_thread_state AS
-WITH span AS MATERIALIZED (
-    SELECT * FROM _critical_path
-  ),
-  span_starts AS (
-    SELECT
-      span.id,
-      span.utid,
-      span.critical_path_id,
-      span.critical_path_blocked_dur,
-      span.critical_path_blocked_state,
-      span.critical_path_blocked_function,
-      span.critical_path_utid,
-      thread_state_id,
-      MAX(thread_state.ts, span.ts) AS ts,
-      span.ts + span.dur AS span_end_ts,
-      thread_state.ts + thread_state.dur AS thread_state_end_ts,
-      thread_state.state,
-      thread_state.function,
-      thread_state.cpu,
-      thread_state.io_wait
-    FROM span
-    JOIN _span_graph_thread_state_sp thread_state USING(id)
-  )
-SELECT
   id,
-  thread_state_id,
   ts,
-  MIN(span_end_ts, thread_state_end_ts) - ts AS dur,
+  dur,
   utid,
-  state,
-  function,
-  cpu,
-  io_wait,
   critical_path_id,
   critical_path_blocked_dur,
   critical_path_blocked_state,
   critical_path_blocked_function,
   critical_path_utid
-FROM span_starts
-WHERE MIN(span_end_ts, thread_state_end_ts) - ts > 0;
-
--- |_thread_executing_span_graph| + thread_state + critical_path span joined with
--- |_thread_executing_span_graph| + slice view.
-CREATE VIRTUAL TABLE _critical_path_sp
-USING
-  SPAN_LEFT_JOIN(
-    _critical_path_thread_state PARTITIONED id,
-     _span_graph_slice PARTITIONED id);
-
--- Flattened slices span joined with their thread_states. This contains the 'self' information
--- without 'critical_path' (blocking) information.
-CREATE VIRTUAL TABLE _self_sp USING
-  SPAN_LEFT_JOIN(thread_state PARTITIONED utid, _slice_flattened PARTITIONED utid);
-
--- Limited view of |_self_sp|.
-CREATE PERFETTO VIEW _self_view
-  AS
-  SELECT
-    id AS self_thread_state_id,
-    slice_id AS self_slice_id,
-    ts,
-    dur,
-    utid AS critical_path_utid,
-    state AS self_state,
-    blocked_function AS self_function,
-    cpu AS self_cpu,
-    io_wait AS self_io_wait,
-    name AS self_slice_name,
-    depth AS self_slice_depth
-    FROM _self_sp;
-
--- Self and critical path span join. This contains the union of the time intervals from the following:
---  a. Self slice stack + thread_state.
---  b. Critical path stack + thread_state.
-CREATE VIRTUAL TABLE _self_and_critical_path_sp
-USING
-  SPAN_JOIN(
-    _self_view PARTITIONED critical_path_utid,
-    _critical_path_sp PARTITIONED critical_path_utid);
-
--- Returns a view of |_self_and_critical_path_sp| unpivoted over the following columns:
--- self thread_state.
--- self blocked_function (if one exists).
--- self process_name (enabled with |enable_process_name|).
--- self thread_name (enabled with |enable_thread_name|).
--- self slice_stack (enabled with |enable_self_slice|).
--- critical_path thread_state.
--- critical_path process_name.
--- critical_path thread_name.
--- critical_path slice_stack (enabled with |enable_critical_path_slice|).
--- running cpu (if one exists).
--- A 'stack' is the group of resulting unpivoted rows sharing the same timestamp.
-CREATE PERFETTO FUNCTION _critical_path_stack(critical_path_utid INT, ts LONG, dur LONG, enable_process_name INT, enable_thread_name INT, enable_self_slice INT, enable_critical_path_slice INT)
-RETURNS
-  TABLE(
-    id INT,
-    ts LONG,
-    dur LONG,
-    utid INT,
-    stack_depth INT,
-    name STRING,
-    table_name STRING,
-    critical_path_utid INT) AS
-  -- Spans filtered to the query time window and critical_path_utid.
-  -- This is a preliminary step that gets the start and end ts of all the rows
-  -- so that we can chop the ends of each interval correctly if it overlaps with the query time interval.
-  WITH relevant_spans_starts AS (
-    SELECT
-      self_thread_state_id,
-      self_state,
-      self_slice_id,
-      self_slice_name,
-      self_slice_depth,
-      self_function,
-      self_io_wait,
-      thread_state_id,
-      state,
-      function,
-      io_wait,
-      slice_id,
-      slice_name,
-      slice_depth,
-      cpu,
-      utid,
-      MAX(ts, $ts) AS ts,
-      MIN(ts + dur, $ts + $dur) AS end_ts,
-      critical_path_utid
-    FROM _self_and_critical_path_sp
-    WHERE dur > 0 AND critical_path_utid = $critical_path_utid
-  ),
-  -- This is the final step that gets the |dur| of each span from the start and
-  -- and end ts of the previous step.
-  -- Now we manually unpivot the result with 3 key steps: 1) Self 2) Critical path 3) CPU
-  -- This CTE is heavily used throughout the entire function so materializing it is
-  -- very important.
-  relevant_spans AS MATERIALIZED (
-    SELECT
-      self_thread_state_id,
-      self_state,
-      self_slice_id,
-      self_slice_name,
-      self_slice_depth,
-      self_function,
-      self_io_wait,
-      thread_state_id,
-      state,
-      function,
-      io_wait,
-      slice_id,
-      slice_name,
-      slice_depth,
-      cpu,
-      utid,
-      ts,
-      end_ts - ts AS dur,
-      critical_path_utid,
-      utid
-    FROM relevant_spans_starts
-    WHERE dur > 0
-  ),
-  -- 1. Builds the 'self' stack of items as an ordered UNION ALL
-  self_stack AS MATERIALIZED (
-    -- Builds the self thread_state
-    SELECT
-      self_thread_state_id AS id,
-      ts,
-      dur,
-      utid,
-      0 AS stack_depth,
-      'thread_state: ' || self_state AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM relevant_spans
-    UNION ALL
-    -- Builds the self kernel blocked_function
-    SELECT
-      self_thread_state_id AS id,
-      ts,
-      dur,
-      utid,
-      1 AS stack_depth,
-      IIF(self_state GLOB 'R*', NULL, 'kernel function: ' || self_function) AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM relevant_spans
-    UNION ALL
-    -- Builds the self kernel io_wait
-    SELECT
-      self_thread_state_id AS id,
-      ts,
-      dur,
-      utid,
-      2 AS stack_depth,
-      IIF(self_state GLOB 'R*', NULL, 'io_wait: ' || self_io_wait) AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM relevant_spans
-    UNION ALL
-    -- Builds the self process_name
-    SELECT
-      self_thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      3 AS stack_depth,
-      IIF($enable_process_name, 'process_name: ' || process.name, NULL) AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM relevant_spans
-    LEFT JOIN thread
-      ON thread.utid = critical_path_utid
-    LEFT JOIN process
-      USING (upid)
-    -- Builds the self thread_name
-    UNION ALL
-    SELECT
-      self_thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      4 AS stack_depth,
-      IIF($enable_thread_name, 'thread_name: ' || thread.name, NULL) AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM relevant_spans
-    LEFT JOIN thread
-      ON thread.utid = critical_path_utid
-    JOIN process
-      USING (upid)
-    UNION ALL
-    -- Builds the self 'ancestor' slice stack
-    SELECT
-      anc.id,
-      slice.ts,
-      slice.dur,
-      slice.utid,
-      anc.depth + 5 AS stack_depth,
-      IIF($enable_self_slice, anc.name, NULL) AS name,
-      'slice' AS table_name,
-      critical_path_utid
-    FROM relevant_spans slice
-    JOIN ancestor_slice(self_slice_id) anc WHERE anc.dur != -1
-    UNION ALL
-    -- Builds the self 'deepest' ancestor slice stack
-    SELECT
-      self_slice_id AS id,
-      ts,
-      dur,
-      utid,
-      self_slice_depth + 5 AS stack_depth,
-      IIF($enable_self_slice, self_slice_name, NULL) AS name,
-      'slice' AS table_name,
-      critical_path_utid
-    FROM relevant_spans slice
-    -- Ordering by stack depth is important to ensure the items can
-    -- be renedered in the UI as a debug track in the order in which
-    -- the sub-queries were 'unioned'.
-    ORDER BY stack_depth
-  ),
-  -- Prepares for stage 2 in building the entire stack.
-  -- Computes the starting depth for each stack. This is necessary because
-  -- each self slice stack has variable depth and the depth in each stack
-  -- most be contiguous in order to efficiently generate a pprof in the future.
-  critical_path_start_depth AS MATERIALIZED (
-    SELECT critical_path_utid, ts, MAX(stack_depth) + 1 AS start_depth
-    FROM self_stack
-    GROUP BY critical_path_utid, ts
-  ),
-  critical_path_span AS MATERIALIZED (
-    SELECT
-      thread_state_id,
-      state,
-      function,
-      io_wait,
-      slice_id,
-      slice_name,
-      slice_depth,
-      spans.ts,
-      spans.dur,
-      spans.critical_path_utid,
-      utid,
-      start_depth
-    FROM relevant_spans spans
-    JOIN critical_path_start_depth
-      ON
-        critical_path_start_depth.critical_path_utid = spans.critical_path_utid
-        AND critical_path_start_depth.ts = spans.ts
-    WHERE critical_path_start_depth.critical_path_utid = $critical_path_utid AND spans.critical_path_utid != spans.utid
-  ),
-  -- 2. Builds the 'critical_path' stack of items as an ordered UNION ALL
-  critical_path_stack AS MATERIALIZED (
-    -- Builds the critical_path thread_state
-    SELECT
-      thread_state_id AS id,
-      ts,
-      dur,
-      utid,
-      start_depth AS stack_depth,
-      'blocking thread_state: ' || state AS name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM critical_path_span
-    UNION ALL
-    -- Builds the critical_path process_name
-    SELECT
-      thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      start_depth + 1 AS stack_depth,
-      'blocking process_name: ' || process.name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM critical_path_span
-    JOIN thread USING (utid)
-    LEFT JOIN process USING (upid)
-    UNION ALL
-    -- Builds the critical_path thread_name
-    SELECT
-      thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      start_depth + 2 AS stack_depth,
-      'blocking thread_name: ' || thread.name,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM critical_path_span
-    JOIN thread USING (utid)
-    UNION ALL
-    -- Builds the critical_path kernel blocked_function
-    SELECT
-      thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      start_depth + 3 AS stack_depth,
-      'blocking kernel_function: ' || function,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM critical_path_span
-    JOIN thread USING (utid)
-    UNION ALL
-    -- Builds the critical_path kernel io_wait
-    SELECT
-      thread_state_id AS id,
-      ts,
-      dur,
-      thread.utid,
-      start_depth + 4 AS stack_depth,
-      'blocking io_wait: ' || io_wait,
-      'thread_state' AS table_name,
-      critical_path_utid
-    FROM critical_path_span
-    JOIN thread USING (utid)
-    UNION ALL
-    -- Builds the critical_path 'ancestor' slice stack
-    SELECT
-      anc.id,
-      slice.ts,
-      slice.dur,
-      slice.utid,
-      anc.depth + start_depth + 5 AS stack_depth,
-      IIF($enable_critical_path_slice, anc.name, NULL) AS name,
-      'slice' AS table_name,
-      critical_path_utid
-    FROM critical_path_span slice
-    JOIN ancestor_slice(slice_id) anc WHERE anc.dur != -1
-    UNION ALL
-    -- Builds the critical_path 'deepest' slice
-    SELECT
-      slice_id AS id,
-      ts,
-      dur,
-      utid,
-      slice_depth + start_depth + 5 AS stack_depth,
-      IIF($enable_critical_path_slice, slice_name, NULL) AS name,
-      'slice' AS table_name,
-      critical_path_utid
-    FROM critical_path_span slice
-    -- Ordering is also important as in the 'self' step above.
-    ORDER BY stack_depth
-  ),
-  -- Prepares for stage 3 in building the entire stack.
-  -- Computes the starting depth for each stack using the deepest stack_depth between
-  -- the critical_path stack and self stack. The self stack depth is
-  -- already computed and materialized in |critical_path_start_depth|.
-  cpu_start_depth_raw AS (
-    SELECT critical_path_utid, ts, MAX(stack_depth) + 1 AS start_depth
-    FROM critical_path_stack
-    GROUP BY critical_path_utid, ts
-    UNION ALL
-    SELECT * FROM critical_path_start_depth
-  ),
-  cpu_start_depth AS (
-    SELECT critical_path_utid, ts, MAX(start_depth) AS start_depth
-    FROM cpu_start_depth_raw
-    GROUP BY critical_path_utid, ts
-  ),
-  -- 3. Builds the 'CPU' stack for 'Running' states in either the self or critical path stack.
-  cpu_stack AS (
-    SELECT
-      thread_state_id AS id,
-      spans.ts,
-      spans.dur,
-      utid,
-      start_depth AS stack_depth,
-      'cpu: ' || cpu AS name,
-      'thread_state' AS table_name,
-      spans.critical_path_utid
-    FROM relevant_spans spans
-    JOIN cpu_start_depth
-      ON
-        cpu_start_depth.critical_path_utid = spans.critical_path_utid
-        AND cpu_start_depth.ts = spans.ts
-    WHERE cpu_start_depth.critical_path_utid = $critical_path_utid AND state = 'Running' OR self_state = 'Running'
-  ),
-  merged AS (
-    SELECT * FROM self_stack
-    UNION ALL
-    SELECT * FROM critical_path_stack
-    UNION ALL
-    SELECT * FROM cpu_stack
-  )
-SELECT * FROM merged WHERE id IS NOT NULL;
-
--- Critical path stack of thread_executing_spans with the following entities in the critical path
--- stacked from top to bottom: self thread_state, self blocked_function, self process_name,
--- self thread_name, slice stack, critical_path thread_state, critical_path process_name,
--- critical_path thread_name, critical_path slice_stack, running_cpu.
-CREATE PERFETTO FUNCTION _thread_executing_span_critical_path_stack(
-  -- Thread utid to filter critical paths to.
-  critical_path_utid INT,
-  -- Timestamp of start of time range to filter critical paths to.
-  ts LONG,
-  -- Duration of time range to filter critical paths to.
-  dur LONG)
-RETURNS
-  TABLE(
-    -- Id of the thread_state or slice in the thread_executing_span.
-    id INT,
-    -- Timestamp of slice in the critical path.
-    ts LONG,
-    -- Duration of slice in the critical path.
-    dur LONG,
-    -- Utid of thread that emitted the slice.
-    utid INT,
-    -- Stack depth of the entitity in the debug track.
-    stack_depth INT,
-    -- Name of entity in the critical path (could be a thread_state, kernel blocked_function, process_name, thread_name, slice name or cpu).
-    name STRING,
-    -- Table name of entity in the critical path (could be either slice or thread_state).
-    table_name STRING,
-    -- Utid of the thread the critical path was filtered to.
-    critical_path_utid INT
-) AS
-SELECT * FROM _critical_path_stack($critical_path_utid, $ts, $dur, 1, 1, 1, 1);
-
--- Returns a pprof aggregation of the stacks in |_critical_path_stack|.
-CREATE PERFETTO FUNCTION _critical_path_graph(graph_title STRING, critical_path_utid INT, ts LONG, dur LONG, enable_process_name INT, enable_thread_name INT, enable_self_slice INT, enable_critical_path_slice INT)
-RETURNS TABLE(pprof BYTES)
-AS
-WITH
-  stack AS MATERIALIZED (
-    SELECT
-      ts,
-      dur - IFNULL(LEAD(dur) OVER (PARTITION BY critical_path_utid, ts ORDER BY stack_depth), 0) AS dur,
-      name,
-      utid,
-      critical_path_utid,
-      stack_depth
-    FROM
-      _critical_path_stack($critical_path_utid, $ts, $dur, $enable_process_name, $enable_thread_name, $enable_self_slice, $enable_critical_path_slice)
-  ),
-  graph AS (
-    SELECT CAT_STACKS($graph_title) AS stack
-  ),
-  parent AS (
-    SELECT
-      cr.ts,
-      cr.dur,
-      cr.name,
-      cr.utid,
-      cr.stack_depth,
-      CAT_STACKS(graph.stack, cr.name) AS stack,
-      cr.critical_path_utid
-    FROM stack cr, graph
-    WHERE stack_depth = 0
-    UNION ALL
-    SELECT
-      child.ts,
-      child.dur,
-      child.name,
-      child.utid,
-      child.stack_depth,
-      CAT_STACKS(stack, child.name) AS stack,
-      child.critical_path_utid
-    FROM stack child
-    JOIN parent
-      ON
-        parent.critical_path_utid = child.critical_path_utid
-        AND parent.ts = child.ts
-        AND child.stack_depth = parent.stack_depth + 1
-  ),
-  stacks AS (
-    SELECT dur, stack FROM parent
-  )
-SELECT EXPERIMENTAL_PROFILE(stack, 'duration', 'ns', dur) AS pprof FROM stacks;
-
--- Returns a pprof aggreagation of the stacks in |_thread_executing_span_critical_path_stack|
-CREATE PERFETTO FUNCTION _thread_executing_span_critical_path_graph(
-  -- Descriptive name for the graph.
-  graph_title STRING,
-  -- Thread utid to filter critical paths to.
-  critical_path_utid INT,
-  -- Timestamp of start of time range to filter critical paths to.
-  ts INT,
-  -- Duration of time range to filter critical paths to.
-  dur INT)
-RETURNS TABLE(
-  -- Pprof of critical path stacks.
-  pprof BYTES
-)
-AS
-SELECT * FROM _critical_path_graph($graph_title, $critical_path_utid, $ts, $dur, 1, 1, 1, 1);
+FROM _critical_path_by_roots!((SELECT id FROM _wakeup_graph));
diff --git a/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span_with_slice.sql b/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span_with_slice.sql
new file mode 100644
index 0000000..89085f3
--- /dev/null
+++ b/src/trace_processor/perfetto_sql/stdlib/sched/thread_executing_span_with_slice.sql
@@ -0,0 +1,573 @@
+--
+-- Copyright 2024 The Android Open Source Project
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--     https://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+
+INCLUDE PERFETTO MODULE slices.flat_slices;
+INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+CREATE PERFETTO TABLE _critical_path_all AS
+SELECT * FROM  _thread_executing_span_critical_path_all();
+
+-- Limited thread_state view that will later be span joined with the |_thread_executing_span_graph|.
+CREATE PERFETTO VIEW _span_thread_state_view
+AS SELECT id AS thread_state_id, ts, dur, utid, state, blocked_function as function, io_wait, cpu FROM thread_state;
+
+-- Limited slice_view that will later be span joined with the |_thread_executing_span_graph|.
+CREATE PERFETTO VIEW _span_slice_view
+AS
+SELECT
+  slice_id,
+  depth AS slice_depth,
+  name AS slice_name,
+  CAST(ts AS INT) AS ts,
+  CAST(dur AS INT) AS dur,
+  utid
+FROM _slice_flattened;
+
+CREATE VIRTUAL TABLE _span_thread_state_slice_view
+USING
+  SPAN_LEFT_JOIN(
+    _span_thread_state_view PARTITIONED utid,
+    _span_slice_view PARTITIONED utid);
+
+-- |_thread_executing_span_graph| span joined with thread_state information.
+CREATE VIRTUAL TABLE _span_critical_path_thread_state_slice_sp
+USING
+  SPAN_JOIN(
+    _critical_path_all PARTITIONED utid,
+    _span_thread_state_slice_view PARTITIONED utid);
+
+-- |_thread_executing_span_graph| + thread_state view joined with critical_path information.
+CREATE PERFETTO TABLE _critical_path_thread_state_slice AS
+WITH span_starts AS (
+    SELECT
+      span.id,
+      span.utid,
+      span.critical_path_id,
+      span.critical_path_blocked_dur,
+      span.critical_path_blocked_state,
+      span.critical_path_blocked_function,
+      span.critical_path_utid,
+      thread_state_id,
+      MAX(thread_state.ts, span.ts) AS ts,
+      span.ts + span.dur AS span_end_ts,
+      thread_state.ts + thread_state.dur AS thread_state_end_ts,
+      thread_state.state,
+      thread_state.function,
+      thread_state.cpu,
+      thread_state.io_wait,
+      thread_state.slice_id,
+      thread_state.slice_name,
+      thread_state.slice_depth
+    FROM _critical_path_all span
+    JOIN _span_critical_path_thread_state_slice_sp thread_state USING(id)
+  )
+SELECT
+  id,
+  thread_state_id,
+  ts,
+  MIN(span_end_ts, thread_state_end_ts) - ts AS dur,
+  utid,
+  state,
+  function,
+  cpu,
+  io_wait,
+  slice_id,
+  slice_name,
+  slice_depth,
+  critical_path_id,
+  critical_path_blocked_dur,
+  critical_path_blocked_state,
+  critical_path_blocked_function,
+  critical_path_utid
+FROM span_starts
+WHERE MIN(span_end_ts, thread_state_end_ts) - ts > 0;
+
+-- Flattened slices span joined with their thread_states. This contains the 'self' information
+-- without 'critical_path' (blocking) information.
+CREATE VIRTUAL TABLE _self_sp USING
+  SPAN_LEFT_JOIN(thread_state PARTITIONED utid, _slice_flattened PARTITIONED utid);
+
+-- Limited view of |_self_sp|.
+CREATE PERFETTO VIEW _self_view
+  AS
+  SELECT
+    id AS self_thread_state_id,
+    slice_id AS self_slice_id,
+    ts,
+    dur,
+    utid AS critical_path_utid,
+    state AS self_state,
+    blocked_function AS self_function,
+    cpu AS self_cpu,
+    io_wait AS self_io_wait,
+    name AS self_slice_name,
+    depth AS self_slice_depth
+    FROM _self_sp;
+
+-- Self and critical path span join. This contains the union of the time intervals from the following:
+--  a. Self slice stack + thread_state.
+--  b. Critical path stack + thread_state.
+CREATE VIRTUAL TABLE _self_and_critical_path_sp
+USING
+  SPAN_JOIN(
+    _self_view PARTITIONED critical_path_utid,
+    _critical_path_thread_state_slice PARTITIONED critical_path_utid);
+
+-- Returns a view of |_self_and_critical_path_sp| unpivoted over the following columns:
+-- self thread_state.
+-- self blocked_function (if one exists).
+-- self process_name (enabled with |enable_process_name|).
+-- self thread_name (enabled with |enable_thread_name|).
+-- self slice_stack (enabled with |enable_self_slice|).
+-- critical_path thread_state.
+-- critical_path process_name.
+-- critical_path thread_name.
+-- critical_path slice_stack (enabled with |enable_critical_path_slice|).
+-- running cpu (if one exists).
+-- A 'stack' is the group of resulting unpivoted rows sharing the same timestamp.
+CREATE PERFETTO FUNCTION _critical_path_stack(critical_path_utid INT, ts LONG, dur LONG, enable_process_name INT, enable_thread_name INT, enable_self_slice INT, enable_critical_path_slice INT)
+RETURNS
+  TABLE(
+    id INT,
+    ts LONG,
+    dur LONG,
+    utid INT,
+    stack_depth INT,
+    name STRING,
+    table_name STRING,
+    critical_path_utid INT) AS
+  -- Spans filtered to the query time window and critical_path_utid.
+  -- This is a preliminary step that gets the start and end ts of all the rows
+  -- so that we can chop the ends of each interval correctly if it overlaps with the query time interval.
+  WITH relevant_spans_starts AS (
+    SELECT
+      self_thread_state_id,
+      self_state,
+      self_slice_id,
+      self_slice_name,
+      self_slice_depth,
+      self_function,
+      self_io_wait,
+      thread_state_id,
+      state,
+      function,
+      io_wait,
+      slice_id,
+      slice_name,
+      slice_depth,
+      cpu,
+      utid,
+      MAX(ts, $ts) AS ts,
+      MIN(ts + dur, $ts + $dur) AS end_ts,
+      critical_path_utid
+    FROM _self_and_critical_path_sp
+    WHERE dur > 0 AND critical_path_utid = $critical_path_utid
+  ),
+  -- This is the final step that gets the |dur| of each span from the start and
+  -- and end ts of the previous step.
+  -- Now we manually unpivot the result with 3 key steps: 1) Self 2) Critical path 3) CPU
+  -- This CTE is heavily used throughout the entire function so materializing it is
+  -- very important.
+  relevant_spans AS MATERIALIZED (
+    SELECT
+      self_thread_state_id,
+      self_state,
+      self_slice_id,
+      self_slice_name,
+      self_slice_depth,
+      self_function,
+      self_io_wait,
+      thread_state_id,
+      state,
+      function,
+      io_wait,
+      slice_id,
+      slice_name,
+      slice_depth,
+      cpu,
+      utid,
+      ts,
+      end_ts - ts AS dur,
+      critical_path_utid,
+      utid
+    FROM relevant_spans_starts
+    WHERE dur > 0
+  ),
+  -- 1. Builds the 'self' stack of items as an ordered UNION ALL
+  self_stack AS MATERIALIZED (
+    -- Builds the self thread_state
+    SELECT
+      self_thread_state_id AS id,
+      ts,
+      dur,
+      utid,
+      0 AS stack_depth,
+      'thread_state: ' || self_state AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM relevant_spans
+    UNION ALL
+    -- Builds the self kernel blocked_function
+    SELECT
+      self_thread_state_id AS id,
+      ts,
+      dur,
+      utid,
+      1 AS stack_depth,
+      IIF(self_state GLOB 'R*', NULL, 'kernel function: ' || self_function) AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM relevant_spans
+    UNION ALL
+    -- Builds the self kernel io_wait
+    SELECT
+      self_thread_state_id AS id,
+      ts,
+      dur,
+      utid,
+      2 AS stack_depth,
+      IIF(self_state GLOB 'R*', NULL, 'io_wait: ' || self_io_wait) AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM relevant_spans
+    UNION ALL
+    -- Builds the self process_name
+    SELECT
+      self_thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      3 AS stack_depth,
+      IIF($enable_process_name, 'process_name: ' || process.name, NULL) AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM relevant_spans
+    LEFT JOIN thread
+      ON thread.utid = critical_path_utid
+    LEFT JOIN process
+      USING (upid)
+    -- Builds the self thread_name
+    UNION ALL
+    SELECT
+      self_thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      4 AS stack_depth,
+      IIF($enable_thread_name, 'thread_name: ' || thread.name, NULL) AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM relevant_spans
+    LEFT JOIN thread
+      ON thread.utid = critical_path_utid
+    JOIN process
+      USING (upid)
+    UNION ALL
+    -- Builds the self 'ancestor' slice stack
+    SELECT
+      anc.id,
+      slice.ts,
+      slice.dur,
+      slice.utid,
+      anc.depth + 5 AS stack_depth,
+      IIF($enable_self_slice, anc.name, NULL) AS name,
+      'slice' AS table_name,
+      critical_path_utid
+    FROM relevant_spans slice
+    JOIN ancestor_slice(self_slice_id) anc WHERE anc.dur != -1
+    UNION ALL
+    -- Builds the self 'deepest' ancestor slice stack
+    SELECT
+      self_slice_id AS id,
+      ts,
+      dur,
+      utid,
+      self_slice_depth + 5 AS stack_depth,
+      IIF($enable_self_slice, self_slice_name, NULL) AS name,
+      'slice' AS table_name,
+      critical_path_utid
+    FROM relevant_spans slice
+    -- Ordering by stack depth is important to ensure the items can
+    -- be renedered in the UI as a debug track in the order in which
+    -- the sub-queries were 'unioned'.
+    ORDER BY stack_depth
+  ),
+  -- Prepares for stage 2 in building the entire stack.
+  -- Computes the starting depth for each stack. This is necessary because
+  -- each self slice stack has variable depth and the depth in each stack
+  -- most be contiguous in order to efficiently generate a pprof in the future.
+  critical_path_start_depth AS MATERIALIZED (
+    SELECT critical_path_utid, ts, MAX(stack_depth) + 1 AS start_depth
+    FROM self_stack
+    GROUP BY critical_path_utid, ts
+  ),
+  critical_path_span AS MATERIALIZED (
+    SELECT
+      thread_state_id,
+      state,
+      function,
+      io_wait,
+      slice_id,
+      slice_name,
+      slice_depth,
+      spans.ts,
+      spans.dur,
+      spans.critical_path_utid,
+      utid,
+      start_depth
+    FROM relevant_spans spans
+    JOIN critical_path_start_depth
+      ON
+        critical_path_start_depth.critical_path_utid = spans.critical_path_utid
+        AND critical_path_start_depth.ts = spans.ts
+    WHERE critical_path_start_depth.critical_path_utid = $critical_path_utid AND spans.critical_path_utid != spans.utid
+  ),
+  -- 2. Builds the 'critical_path' stack of items as an ordered UNION ALL
+  critical_path_stack AS MATERIALIZED (
+    -- Builds the critical_path thread_state
+    SELECT
+      thread_state_id AS id,
+      ts,
+      dur,
+      utid,
+      start_depth AS stack_depth,
+      'blocking thread_state: ' || state AS name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM critical_path_span
+    UNION ALL
+    -- Builds the critical_path process_name
+    SELECT
+      thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      start_depth + 1 AS stack_depth,
+      'blocking process_name: ' || process.name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM critical_path_span
+    JOIN thread USING (utid)
+    LEFT JOIN process USING (upid)
+    UNION ALL
+    -- Builds the critical_path thread_name
+    SELECT
+      thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      start_depth + 2 AS stack_depth,
+      'blocking thread_name: ' || thread.name,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM critical_path_span
+    JOIN thread USING (utid)
+    UNION ALL
+    -- Builds the critical_path kernel blocked_function
+    SELECT
+      thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      start_depth + 3 AS stack_depth,
+      'blocking kernel_function: ' || function,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM critical_path_span
+    JOIN thread USING (utid)
+    UNION ALL
+    -- Builds the critical_path kernel io_wait
+    SELECT
+      thread_state_id AS id,
+      ts,
+      dur,
+      thread.utid,
+      start_depth + 4 AS stack_depth,
+      'blocking io_wait: ' || io_wait,
+      'thread_state' AS table_name,
+      critical_path_utid
+    FROM critical_path_span
+    JOIN thread USING (utid)
+    UNION ALL
+    -- Builds the critical_path 'ancestor' slice stack
+    SELECT
+      anc.id,
+      slice.ts,
+      slice.dur,
+      slice.utid,
+      anc.depth + start_depth + 5 AS stack_depth,
+      IIF($enable_critical_path_slice, anc.name, NULL) AS name,
+      'slice' AS table_name,
+      critical_path_utid
+    FROM critical_path_span slice
+    JOIN ancestor_slice(slice_id) anc WHERE anc.dur != -1
+    UNION ALL
+    -- Builds the critical_path 'deepest' slice
+    SELECT
+      slice_id AS id,
+      ts,
+      dur,
+      utid,
+      slice_depth + start_depth + 5 AS stack_depth,
+      IIF($enable_critical_path_slice, slice_name, NULL) AS name,
+      'slice' AS table_name,
+      critical_path_utid
+    FROM critical_path_span slice
+    -- Ordering is also important as in the 'self' step above.
+    ORDER BY stack_depth
+  ),
+  -- Prepares for stage 3 in building the entire stack.
+  -- Computes the starting depth for each stack using the deepest stack_depth between
+  -- the critical_path stack and self stack. The self stack depth is
+  -- already computed and materialized in |critical_path_start_depth|.
+  cpu_start_depth_raw AS (
+    SELECT critical_path_utid, ts, MAX(stack_depth) + 1 AS start_depth
+    FROM critical_path_stack
+    GROUP BY critical_path_utid, ts
+    UNION ALL
+    SELECT * FROM critical_path_start_depth
+  ),
+  cpu_start_depth AS (
+    SELECT critical_path_utid, ts, MAX(start_depth) AS start_depth
+    FROM cpu_start_depth_raw
+    GROUP BY critical_path_utid, ts
+  ),
+  -- 3. Builds the 'CPU' stack for 'Running' states in either the self or critical path stack.
+  cpu_stack AS (
+    SELECT
+      thread_state_id AS id,
+      spans.ts,
+      spans.dur,
+      utid,
+      start_depth AS stack_depth,
+      'cpu: ' || cpu AS name,
+      'thread_state' AS table_name,
+      spans.critical_path_utid
+    FROM relevant_spans spans
+    JOIN cpu_start_depth
+      ON
+        cpu_start_depth.critical_path_utid = spans.critical_path_utid
+        AND cpu_start_depth.ts = spans.ts
+    WHERE cpu_start_depth.critical_path_utid = $critical_path_utid AND state = 'Running' OR self_state = 'Running'
+  ),
+  merged AS (
+    SELECT * FROM self_stack
+    UNION ALL
+    SELECT * FROM critical_path_stack
+    UNION ALL
+    SELECT * FROM cpu_stack
+  )
+SELECT * FROM merged WHERE id IS NOT NULL;
+
+-- Critical path stack of thread_executing_spans with the following entities in the critical path
+-- stacked from top to bottom: self thread_state, self blocked_function, self process_name,
+-- self thread_name, slice stack, critical_path thread_state, critical_path process_name,
+-- critical_path thread_name, critical_path slice_stack, running_cpu.
+CREATE PERFETTO FUNCTION _thread_executing_span_critical_path_stack(
+  -- Thread utid to filter critical paths to.
+  critical_path_utid INT,
+  -- Timestamp of start of time range to filter critical paths to.
+  ts LONG,
+  -- Duration of time range to filter critical paths to.
+  dur LONG)
+RETURNS
+  TABLE(
+    -- Id of the thread_state or slice in the thread_executing_span.
+    id INT,
+    -- Timestamp of slice in the critical path.
+    ts LONG,
+    -- Duration of slice in the critical path.
+    dur LONG,
+    -- Utid of thread that emitted the slice.
+    utid INT,
+    -- Stack depth of the entitity in the debug track.
+    stack_depth INT,
+    -- Name of entity in the critical path (could be a thread_state, kernel blocked_function, process_name, thread_name, slice name or cpu).
+    name STRING,
+    -- Table name of entity in the critical path (could be either slice or thread_state).
+    table_name STRING,
+    -- Utid of the thread the critical path was filtered to.
+    critical_path_utid INT
+) AS
+SELECT * FROM _critical_path_stack($critical_path_utid, $ts, $dur, 1, 1, 1, 1);
+
+-- Returns a pprof aggregation of the stacks in |_critical_path_stack|.
+CREATE PERFETTO FUNCTION _critical_path_graph(graph_title STRING, critical_path_utid INT, ts LONG, dur LONG, enable_process_name INT, enable_thread_name INT, enable_self_slice INT, enable_critical_path_slice INT)
+RETURNS TABLE(pprof BYTES)
+AS
+WITH
+  stack AS MATERIALIZED (
+    SELECT
+      ts,
+      dur - IFNULL(LEAD(dur) OVER (PARTITION BY critical_path_utid, ts ORDER BY stack_depth), 0) AS dur,
+      name,
+      utid,
+      critical_path_utid,
+      stack_depth
+    FROM
+      _critical_path_stack($critical_path_utid, $ts, $dur, $enable_process_name, $enable_thread_name, $enable_self_slice, $enable_critical_path_slice)
+  ),
+  graph AS (
+    SELECT CAT_STACKS($graph_title) AS stack
+  ),
+  parent AS (
+    SELECT
+      cr.ts,
+      cr.dur,
+      cr.name,
+      cr.utid,
+      cr.stack_depth,
+      CAT_STACKS(graph.stack, cr.name) AS stack,
+      cr.critical_path_utid
+    FROM stack cr, graph
+    WHERE stack_depth = 0
+    UNION ALL
+    SELECT
+      child.ts,
+      child.dur,
+      child.name,
+      child.utid,
+      child.stack_depth,
+      CAT_STACKS(stack, child.name) AS stack,
+      child.critical_path_utid
+    FROM stack child
+    JOIN parent
+      ON
+        parent.critical_path_utid = child.critical_path_utid
+        AND parent.ts = child.ts
+        AND child.stack_depth = parent.stack_depth + 1
+  ),
+  stacks AS (
+    SELECT dur, stack FROM parent
+  )
+SELECT EXPERIMENTAL_PROFILE(stack, 'duration', 'ns', dur) AS pprof FROM stacks;
+
+-- Returns a pprof aggreagation of the stacks in |_thread_executing_span_critical_path_stack|
+CREATE PERFETTO FUNCTION _thread_executing_span_critical_path_graph(
+  -- Descriptive name for the graph.
+  graph_title STRING,
+  -- Thread utid to filter critical paths to.
+  critical_path_utid INT,
+  -- Timestamp of start of time range to filter critical paths to.
+  ts INT,
+  -- Duration of time range to filter critical paths to.
+  dur INT)
+RETURNS TABLE(
+  -- Pprof of critical path stacks.
+  pprof BYTES
+)
+AS
+SELECT * FROM _critical_path_graph($graph_title, $critical_path_utid, $ts, $dur, 1, 1, 1, 1);
diff --git a/test/trace_processor/diff_tests/tables/tests_sched.py b/test/trace_processor/diff_tests/tables/tests_sched.py
index 71ca6ab..c2ad002 100644
--- a/test/trace_processor/diff_tests/tables/tests_sched.py
+++ b/test/trace_processor/diff_tests/tables/tests_sched.py
@@ -146,34 +146,32 @@
         query="""
         INCLUDE PERFETTO MODULE sched.thread_executing_span;
         SELECT
-          root_id,
-          parent_id,
+          waker_id,
+          prev_id,
+          prev_end_ts,
           id,
           ts,
-          dur,
+          end_ts,
+          is_kernel,
           utid,
-          blocked_dur,
-          blocked_state,
-          blocked_function,
-          is_root,
-          depth
-        FROM _thread_executing_span_graph
-          WHERE blocked_function IS NOT NULL
+          state,
+          blocked_function
+        FROM _wakeup_graph
         ORDER BY ts
         LIMIT 10
         """,
         out=Csv("""
-        "root_id","parent_id","id","ts","dur","utid","blocked_dur","blocked_state","blocked_function","is_root","depth"
-        357,377,380,1735842234188,283571,46,351402620,"I","worker_thread",0,5
-        394,402,405,1735843726296,8545303,46,1208537,"I","worker_thread",0,3
-        357,419,432,1735850643698,16245,95,154087,"I","worker_thread",0,4
-        357,443,446,1735851953029,554638012,95,1103252,"I","worker_thread",0,6
-        357,500,503,1735886367018,191863,46,34095419,"I","worker_thread",0,10
-        357,446,667,1736125372478,52493,46,238813597,"I","worker_thread",0,7
-        357,835,838,1736405409972,278036,46,279985001,"I","worker_thread",0,12
-        357,862,865,1736406817672,7959441,46,1129664,"I","worker_thread",0,10
-        357,882,889,1736413734042,25870,95,7143001,"I","worker_thread",0,11
-        357,882,894,1736413763072,31692550,11,4413060,"I","rcu_gp_fqs_loop",0,11
+        "waker_id","prev_id","prev_end_ts","id","ts","end_ts","is_kernel","utid","state","blocked_function"
+        "[NULL]","[NULL]","[NULL]",5,1735489812571,1735489896509,0,304,"[NULL]","[NULL]"
+        6,"[NULL]","[NULL]",11,1735489876788,1735489953773,0,428,"[NULL]","[NULL]"
+        5,"[NULL]","[NULL]",12,1735489879097,1735490217277,0,243,"[NULL]","[NULL]"
+        11,"[NULL]","[NULL]",17,1735489933912,1735490587658,0,230,"[NULL]","[NULL]"
+        "[NULL]","[NULL]","[NULL]",20,1735489972385,1735489995809,0,298,"[NULL]","[NULL]"
+        "[NULL]",20,1735489995809,25,1735489999987,1735490055966,0,298,"S","[NULL]"
+        25,"[NULL]","[NULL]",28,1735490039439,1735490610238,0,421,"[NULL]","[NULL]"
+        25,"[NULL]","[NULL]",29,1735490042084,1735490068213,0,420,"[NULL]","[NULL]"
+        25,"[NULL]","[NULL]",30,1735490045825,1735491418790,0,1,"[NULL]","[NULL]"
+        17,"[NULL]","[NULL]",41,1735490544063,1735490598211,0,427,"[NULL]","[NULL]"
         """))
 
   def test_thread_executing_span_graph_contains_forked_states(self):
@@ -182,23 +180,15 @@
         query="""
         INCLUDE PERFETTO MODULE sched.thread_executing_span;
         SELECT
-          root_id,
-          parent_id,
           id,
-          ts,
-          dur,
-          utid,
-          blocked_dur,
-          blocked_state,
-          blocked_function,
-          is_root,
-          depth
-        FROM _thread_executing_span_graph
-          WHERE ts = 1735842081507 AND dur = 293868
+          waker_id,
+          prev_id
+        FROM _wakeup_graph
+          WHERE ts = 1735842081507 AND end_ts = 1735842081507 + 293868
         """,
         out=Csv("""
-        "root_id","parent_id","id","ts","dur","utid","blocked_dur","blocked_state","blocked_function","is_root","depth"
-        357,369,376,1735842081507,293868,1465,"[NULL]","[NULL]","[NULL]",0,4
+        "id","waker_id","prev_id"
+        376,369,"[NULL]"
         """))
 
   def test_thread_executing_span_runnable_state_has_no_running(self):
@@ -218,11 +208,11 @@
         trace=DataPath('sched_wakeup_trace.atr'),
         query="""
         INCLUDE PERFETTO MODULE sched.thread_executing_span;
-        SELECT ts,dur FROM _thread_executing_span_graph
-          WHERE dur IS NULL OR ts IS NULL
+        SELECT ts,end_ts FROM _wakeup_graph
+          WHERE end_ts IS NULL OR ts IS NULL
         """,
         out=Csv("""
-        "ts","dur"
+        "ts","end_ts"
         """))
 
   def test_thread_executing_span_graph_accepts_null_irq_context(self):
@@ -230,44 +220,222 @@
         trace=DataPath('sched_switch_original.pb'),
         query="""
         INCLUDE PERFETTO MODULE sched.thread_executing_span;
-        SELECT COUNT(*) AS count FROM _thread_executing_span_graph
+        SELECT COUNT(*) AS count FROM _wakeup_graph
         """,
         out=Csv("""
         "count"
-        25
+        17
         """))
 
-  def test_thread_executing_span_critical_path_all(self):
+  def test_thread_executing_span_flatten_critical_path_tasks(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_switch_original.pb'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        CREATE PERFETTO TABLE graph AS
+        SELECT
+          id AS source_node_id,
+          COALESCE(waker_id, id) AS dest_node_id,
+          id - COALESCE(waker_id, id) AS edge_weight
+        FROM _wakeup_graph;
+
+        CREATE PERFETTO TABLE roots AS
+        SELECT
+          _wakeup_graph.id AS root_node_id,
+          _wakeup_graph.id - COALESCE(prev_id, _wakeup_graph.id) AS root_target_weight,
+          id,
+          ts,
+          end_ts,
+          utid
+        FROM _wakeup_graph LIMIT 10;
+
+        CREATE PERFETTO TABLE critical_path AS
+        SELECT * FROM graph_reachable_weight_bounded_dfs!(graph, roots, 1);
+
+        SELECT * FROM _flatten_critical_path_tasks!(critical_path);
+        """,
+        out=Csv("""
+        "ts","root_node_id","node_id","dur","node_utid","prev_end_ts"
+        807082868359903,29,29,"[NULL]",8,"[NULL]"
+        807082871734539,35,35,"[NULL]",9,"[NULL]"
+        807082871734539,38,35,45052,9,"[NULL]"
+        807082871779591,38,38,"[NULL]",5,807082871764903
+        807082878623081,45,45,"[NULL]",9,807082871805424
+        807082947156994,57,57,"[NULL]",9,807082878865945
+        807082947246838,62,62,"[NULL]",6,807082879179539
+        807082947261525,63,63,"[NULL]",12,"[NULL]"
+        807082947267463,64,64,"[NULL]",13,"[NULL]"
+        807082947278140,65,65,"[NULL]",14,"[NULL]"
+        807082947288765,66,66,"[NULL]",15,"[NULL]"
+        """))
+
+  def test_thread_executing_span_intervals_to_roots_edge_case(self):
     return DiffTestBlueprint(
         trace=DataPath('sched_wakeup_trace.atr'),
         query="""
         INCLUDE PERFETTO MODULE sched.thread_executing_span;
-        SELECT
-          id,
-          ts,
-          dur,
-          utid,
-          critical_path_id,
-          critical_path_blocked_dur,
-          critical_path_blocked_state,
-          critical_path_blocked_function,
-          critical_path_utid INT
-        FROM _thread_executing_span_critical_path(NULL, start_ts, end_ts), trace_bounds
-        ORDER BY ts
-        LIMIT 10
+
+        SELECT * FROM
+        _intervals_to_roots!((SELECT 1477 AS utid, trace_start() AS ts, trace_end() - trace_start() AS dur))
+        LIMIT 10;
         """,
         out=Csv("""
-        "id","ts","dur","utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function","INT"
-        5,1735489812571,83938,304,5,"[NULL]","[NULL]","[NULL]",304
-        6,1735489833977,52463,297,6,"[NULL]","[NULL]","[NULL]",297
-        11,1735489876788,76985,428,11,"[NULL]","[NULL]","[NULL]",428
-        12,1735489879097,338180,243,12,"[NULL]","[NULL]","[NULL]",243
-        17,1735489933912,653746,230,17,"[NULL]","[NULL]","[NULL]",230
-        25,1735489999987,55979,298,25,4178,"S","[NULL]",298
-        28,1735490039439,570799,421,28,"[NULL]","[NULL]","[NULL]",421
-        29,1735490042084,26129,420,29,"[NULL]","[NULL]","[NULL]",420
-        30,1735490045825,1372965,1,30,"[NULL]","[NULL]","[NULL]",1
-        41,1735490544063,54148,427,41,"[NULL]","[NULL]","[NULL]",427
+        "id"
+        11889
+        11892
+        11893
+        11896
+        11897
+        11900
+        11911
+        11916
+        11917
+        11921
+        """))
+
+  def test_thread_executing_span_intervals_to_roots(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_wakeup_trace.atr'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        SELECT * FROM
+        _intervals_to_roots!((SELECT 1477 AS utid, 1737362149192 AS ts, CAST(2e7 AS INT) AS dur))
+        LIMIT 10;
+        """,
+        out=Csv("""
+        "id"
+        11980
+        11983
+        11984
+        11989
+        11990
+        11991
+        11992
+        11993
+        12001
+        12006
+        """))
+
+  def test_thread_executing_span_flatten_critical_paths(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_switch_original.pb'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        CREATE PERFETTO TABLE graph AS
+        SELECT
+          id AS source_node_id,
+          COALESCE(waker_id, id) AS dest_node_id,
+          id - COALESCE(waker_id, id) AS edge_weight
+        FROM _wakeup_graph;
+
+        CREATE PERFETTO TABLE roots AS
+        SELECT
+          _wakeup_graph.id AS root_node_id,
+          _wakeup_graph.id - COALESCE(prev_id, _wakeup_graph.id) AS root_target_weight,
+          id,
+          ts,
+          end_ts,
+          utid
+        FROM _wakeup_graph;
+
+        CREATE PERFETTO TABLE critical_path AS
+        SELECT * FROM graph_reachable_weight_bounded_dfs!(graph, roots, 1);
+
+        SELECT * FROM _flatten_critical_paths!(critical_path, _sleep);
+        """,
+        out=Csv("""
+        "ts","dur","utid","id","root_id","prev_end_ts","critical_path_utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function"
+        807082871764903,14688,9,35,38,"[NULL]",5,38,14688,"S","[NULL]"
+        807082947156994,351302,9,57,76,807082878865945,5,76,68858913,"S","[NULL]"
+        807083031589763,324114,21,127,130,"[NULL]",5,130,80026987,"S","[NULL]"
+        """))
+
+  def test_thread_executing_span_critical_path(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_switch_original.pb'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        CREATE PERFETTO TABLE graph AS
+        SELECT
+          id AS source_node_id,
+          COALESCE(waker_id, id) AS dest_node_id,
+          id - COALESCE(waker_id, id) AS edge_weight
+        FROM _wakeup_graph;
+
+        CREATE PERFETTO TABLE roots AS
+        SELECT
+          _wakeup_graph.id AS root_node_id,
+          _wakeup_graph.id - COALESCE(prev_id, _wakeup_graph.id) AS root_target_weight,
+          id,
+          ts,
+          end_ts,
+          utid
+        FROM _wakeup_graph;
+
+        SELECT * FROM _critical_path!(graph, roots, _sleep);
+        """,
+        out=Csv("""
+        "ts","dur","root_id","id","utid","critical_path_utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function"
+        807082868359903,81302,29,29,8,8,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082871734539,70885,35,35,9,9,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082871764903,14688,38,35,9,5,38,14688,"S","[NULL]"
+        807082871779591,55729,38,38,5,5,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082878623081,242864,45,45,9,9,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947156994,436354,57,57,9,9,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947246838,1038854,62,62,6,6,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947261525,293594,63,63,12,12,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947267463,228958,64,64,13,13,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947278140,54114,65,65,14,14,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947288765,338802,66,66,15,15,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947294182,296875,67,67,16,16,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082947156994,351302,76,57,9,5,76,68858913,"S","[NULL]"
+        807082947508296,122083,76,76,5,5,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082951822463,104427,96,96,9,9,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807082959173506,215104,107,107,6,6,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807083031589763,436198,127,127,21,21,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807083031589763,324114,130,127,21,5,130,80026987,"S","[NULL]"
+        807083031913877,166302,130,130,5,5,"[NULL]","[NULL]","[NULL]","[NULL]"
+        807083032278825,208490,135,135,2,2,"[NULL]","[NULL]","[NULL]","[NULL]"
+        """))
+
+  def test_thread_executing_span_critical_path_by_roots(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_switch_original.pb'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        SELECT * FROM _critical_path_by_roots!(_intervals_to_roots!((SELECT 6 AS utid, trace_start() AS ts, trace_end() - trace_start() AS dur)));
+        """,
+        out=Csv("""
+        "id","ts","dur","utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function","critical_path_utid"
+        62,807082947246838,1038854,6,"[NULL]","[NULL]","[NULL]","[NULL]",6
+        63,807082947261525,293594,12,"[NULL]","[NULL]","[NULL]","[NULL]",12
+        64,807082947267463,228958,13,"[NULL]","[NULL]","[NULL]","[NULL]",13
+        65,807082947278140,54114,14,"[NULL]","[NULL]","[NULL]","[NULL]",14
+        66,807082947288765,338802,15,"[NULL]","[NULL]","[NULL]","[NULL]",15
+        67,807082947294182,296875,16,"[NULL]","[NULL]","[NULL]","[NULL]",16
+        57,807082947156994,351302,9,76,68858913,"S","[NULL]",5
+        76,807082947508296,122083,5,"[NULL]","[NULL]","[NULL]","[NULL]",5
+        96,807082951822463,104427,9,"[NULL]","[NULL]","[NULL]","[NULL]",9
+        107,807082959173506,215104,6,"[NULL]","[NULL]","[NULL]","[NULL]",6
+        """))
+
+  def test_thread_executing_span_critical_path_by_intervals(self):
+    return DiffTestBlueprint(
+        trace=DataPath('sched_switch_original.pb'),
+        query="""
+        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+
+        SELECT * FROM _critical_path_by_intervals!((SELECT 6 AS utid, trace_start() AS ts, trace_end() - trace_start() AS dur));
+        """,
+        out=Csv("""
+        "id","ts","dur","utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function","critical_path_utid"
+        62,807082947246838,1038854,6,"[NULL]","[NULL]","[NULL]","[NULL]",6
+        107,807082959173506,215104,6,"[NULL]","[NULL]","[NULL]","[NULL]",6
         """))
 
   def test_thread_executing_span_critical_path_utid(self):
@@ -284,22 +452,22 @@
           critical_path_blocked_dur,
           critical_path_blocked_state,
           critical_path_blocked_function,
-          critical_path_utid INT
+          critical_path_utid
         FROM _thread_executing_span_critical_path((select utid from thread where tid = 3487), start_ts, end_ts), trace_bounds
         ORDER BY ts
         LIMIT 10
         """,
         out=Csv("""
-        "id","ts","dur","utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function","INT"
-        11889,1737349401439,7705561,1477,11889,"[NULL]","[NULL]","[NULL]",1477
+        "id","ts","dur","utid","critical_path_id","critical_path_blocked_dur","critical_path_blocked_state","critical_path_blocked_function","critical_path_utid"
+        11889,1737349401439,7705561,1477,"[NULL]","[NULL]","[NULL]","[NULL]",1477
         11952,1737357107000,547583,1480,11980,547583,"S","[NULL]",1477
-        11980,1737357654583,8430762,1477,11980,547583,"S","[NULL]",1477
+        11980,1737357654583,8430762,1477,"[NULL]","[NULL]","[NULL]","[NULL]",1477
         12052,1737366085345,50400,91,12057,50400,"S","[NULL]",1477
-        12057,1737366135745,6635927,1477,12057,50400,"S","[NULL]",1477
+        12057,1737366135745,6635927,1477,"[NULL]","[NULL]","[NULL]","[NULL]",1477
         12081,1737372771672,12798314,1488,12254,12798314,"S","[NULL]",1477
-        12254,1737385569986,21830622,1477,12254,12798314,"S","[NULL]",1477
+        12254,1737385569986,21830622,1477,"[NULL]","[NULL]","[NULL]","[NULL]",1477
         12517,1737407400608,241267,91,12521,241267,"S","[NULL]",1477
-        12521,1737407641875,1830015,1477,12521,241267,"S","[NULL]",1477
+        12521,1737407641875,1830015,1477,"[NULL]","[NULL]","[NULL]","[NULL]",1477
         12669,1737409471890,68590,91,12672,68590,"S","[NULL]",1477
         """))
 
@@ -307,7 +475,7 @@
     return DiffTestBlueprint(
         trace=DataPath('sched_wakeup_trace.atr'),
         query="""
-        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+        INCLUDE PERFETTO MODULE sched.thread_executing_span_with_slice;
         SELECT
           id,
           ts,
@@ -340,7 +508,7 @@
     return DiffTestBlueprint(
         trace=DataPath('sched_wakeup_trace.atr'),
         query="""
-        INCLUDE PERFETTO MODULE sched.thread_executing_span;
+        INCLUDE PERFETTO MODULE sched.thread_executing_span_with_slice;
         SELECT HEX(pprof) FROM _thread_executing_span_critical_path_graph("critical path", (select utid from thread where tid = 3487), 1737488133487, 16000), trace_bounds
       """,
         out=BinaryProto(
diff --git a/ui/src/frontend/app.ts b/ui/src/frontend/app.ts
index 940ff31..24769c2 100644
--- a/ui/src/frontend/app.ts
+++ b/ui/src/frontend/app.ts
@@ -322,8 +322,8 @@
 
         if (engine !== undefined && trackUtid != 0) {
           await runQuery(
-            `SELECT IMPORT('sched.thread_executing_span');`,
-            engine,
+              `INCLUDE PERFETTO MODULE sched.thread_executing_span;`,
+              engine,
           );
           await addDebugSliceTrack(
             engine,
@@ -365,8 +365,8 @@
 
         if (engine !== undefined && trackUtid != 0) {
           await runQuery(
-            `SELECT IMPORT('sched.thread_executing_span');`,
-            engine,
+              `INCLUDE PERFETTO MODULE sched.thread_executing_span_with_slice;`,
+              engine,
           );
           await addDebugSliceTrack(
             engine,
@@ -399,7 +399,8 @@
 
         if (engine !== undefined && trackUtid != 0) {
           addQueryResultsTab({
-            query: `SELECT IMPORT('sched.thread_executing_span');
+            query:
+                `INCLUDE PERFETTO MODULE sched.thread_executing_span_with_slice;
                    SELECT *
                       FROM
                         _thread_executing_span_critical_path_graph(
diff --git a/ui/src/frontend/thread_state_tab.ts b/ui/src/frontend/thread_state_tab.ts
index 30209ab..6b0b2ef 100644
--- a/ui/src/frontend/thread_state_tab.ts
+++ b/ui/src/frontend/thread_state_tab.ts
@@ -271,63 +271,62 @@
       `${state.state} for ${renderDuration(state.dur)}`;
     return [
       m(
-        Tree,
-        this.relatedStates.waker &&
-          m(TreeNode, {
+          Tree,
+          this.relatedStates.waker && m(TreeNode, {
             left: 'Waker',
             right: renderRef(
-              this.relatedStates.waker,
-              getFullThreadName(this.relatedStates.waker.thread),
-            ),
+                this.relatedStates.waker,
+                getFullThreadName(this.relatedStates.waker.thread),
+                ),
           }),
-        this.relatedStates.prev &&
-          m(TreeNode, {
+          this.relatedStates.prev && m(TreeNode, {
             left: 'Previous state',
             right: renderRef(
-              this.relatedStates.prev,
-              nameForNextOrPrev(this.relatedStates.prev),
-            ),
+                this.relatedStates.prev,
+                nameForNextOrPrev(this.relatedStates.prev),
+                ),
           }),
-        this.relatedStates.next &&
-          m(TreeNode, {
+          this.relatedStates.next && m(TreeNode, {
             left: 'Next state',
             right: renderRef(
-              this.relatedStates.next,
-              nameForNextOrPrev(this.relatedStates.next),
-            ),
+                this.relatedStates.next,
+                nameForNextOrPrev(this.relatedStates.next),
+                ),
           }),
-        this.relatedStates.wakee &&
-          this.relatedStates.wakee.length > 0 &&
-          m(
-            TreeNode,
-            {
-              left: 'Woken threads',
-            },
-            this.relatedStates.wakee.map((state) =>
-              m(TreeNode, {
-                left: m(Timestamp, {
-                  ts: state.ts,
-                  display: [
-                    'Start+',
-                    m(DurationWidget, {dur: Time.sub(state.ts, startTs)}),
-                  ],
-                }),
-                right: renderRef(state, getFullThreadName(state.thread)),
-              }),
-            ),
+          this.relatedStates.wakee && this.relatedStates.wakee.length > 0 &&
+              m(
+                  TreeNode,
+                  {
+                    left: 'Woken threads',
+                  },
+                  this.relatedStates.wakee.map(
+                      (state) => m(TreeNode, {
+                        left: m(Timestamp, {
+                          ts: state.ts,
+                          display: [
+                            'Start+',
+                            m(DurationWidget,
+                              {dur: Time.sub(state.ts, startTs)}),
+                          ],
+                        }),
+                        right:
+                            renderRef(state, getFullThreadName(state.thread)),
+                      }),
+                      ),
+                  ),
           ),
-      ),
       m(Button, {
         label: 'Critical path lite',
         onclick: () =>
-          runQuery(
-            `INCLUDE PERFETTO MODULE sched.thread_executing_span;`,
-            this.engine,
-          ).then(() =>
-            addDebugSliceTrack(
-              this.engine,
-              {
-                sqlSource: `
+            runQuery(
+                `INCLUDE PERFETTO MODULE sched.thread_executing_span;`,
+                this.engine,
+                )
+                .then(
+                    () => addDebugSliceTrack(
+                        this.engine,
+                        {
+                          sqlSource: `
                     SELECT
                       cr.id,
                       cr.utid,
@@ -345,25 +344,26 @@
                     JOIN thread USING(utid)
                     JOIN process USING(upid)
                   `,
-                columns: sliceLiteColumnNames,
-              },
-              `${this.state?.thread?.name}`,
-              sliceLiteColumns,
-              sliceLiteColumnNames,
-            ),
-          ),
+                          columns: sliceLiteColumnNames,
+                        },
+                        `${this.state?.thread?.name}`,
+                        sliceLiteColumns,
+                        sliceLiteColumnNames,
+                        ),
+                    ),
       }),
       m(Button, {
         label: 'Critical path',
         onclick: () =>
-          runQuery(
-            `INCLUDE PERFETTO MODULE sched.thread_executing_span;`,
-            this.engine,
-          ).then(() =>
-            addDebugSliceTrack(
-              this.engine,
-              {
-                sqlSource: `
+            runQuery(
+                `INCLUDE PERFETTO MODULE sched.thread_executing_span_with_slice;`,
+                this.engine,
+                )
+                .then(
+                    () => addDebugSliceTrack(
+                        this.engine,
+                        {
+                          sqlSource: `
                     SELECT cr.id, cr.utid, cr.ts, cr.dur, cr.name, cr.table_name
                       FROM
                         _thread_executing_span_critical_path_stack(
@@ -372,13 +372,13 @@
                           trace_bounds.end_ts - trace_bounds.start_ts) cr,
                         trace_bounds WHERE name IS NOT NULL
                   `,
-                columns: sliceColumnNames,
-              },
-              `${this.state?.thread?.name}`,
-              sliceColumns,
-              sliceColumnNames,
-            ),
-          ),
+                          columns: sliceColumnNames,
+                        },
+                        `${this.state?.thread?.name}`,
+                        sliceColumns,
+                        sliceColumnNames,
+                        ),
+                    ),
       }),
     ];
   }