diff --git a/source/CalleeFrames.cpp b/source/CalleeFrames.cpp
index 2758a394..a78659ba 100644
--- a/source/CalleeFrames.cpp
+++ b/source/CalleeFrames.cpp
@@ -109,8 +109,13 @@ CalleeFrames CalleeFrames::update_with_propagation_trace(
     // "callee" so we do not need to "propagate" these calls.
     // All these (prior) transform hops are tracked as ExtraTrace hop
     // frames to create a subtrace.
-    // TODO: T158087152 Add extra trace hops for subtraces.
-    return *this;
+    CalleeFrames result = *this;
+    result.map([&propagation_frame](Frame frame) {
+      frame.add_extra_trace(propagation_frame);
+      return frame;
+    });
+
+    return result;
   }
 
   mt_assert(callee_call_info.is_propagation_without_trace());
diff --git a/source/ForwardTaintTransfer.cpp b/source/ForwardTaintTransfer.cpp
index 76e1205b..7dffc091 100644
--- a/source/ForwardTaintTransfer.cpp
+++ b/source/ForwardTaintTransfer.cpp
@@ -288,11 +288,6 @@ void apply_propagations(
       auto transformed_taint_tree =
           transforms::apply_propagation(context, propagation, input_taint_tree);
 
-      if (propagation.call_info().is_propagation_with_trace()) {
-        // Add ExtraTrace first hop
-        transformed_taint_tree.add_extra_trace(propagation);
-      }
-
       auto output_root = propagation_kind->root();
       FeatureMayAlwaysSet features = FeatureMayAlwaysSet::make_always(
           callee.model.add_features_to_arguments(output_root));
diff --git a/source/Taint.cpp b/source/Taint.cpp
index 9a237a20..0a003193 100644
--- a/source/Taint.cpp
+++ b/source/Taint.cpp
@@ -120,13 +120,6 @@ void Taint::add_locally_inferred_features_and_local_position(
   }
 }
 
-void Taint::add_extra_trace(const Frame& propagation_frame) {
-  map([&propagation_frame](Frame frame) {
-    frame.add_extra_trace(propagation_frame);
-    return frame;
-  });
-}
-
 Taint Taint::propagate(
     const Method* callee,
     const AccessPath& callee_port,
diff --git a/source/Taint.h b/source/Taint.h
index f18c2219..4d54d518 100644
--- a/source/Taint.h
+++ b/source/Taint.h
@@ -133,8 +133,6 @@ class Taint final : public sparta::AbstractDomain<Taint> {
       const FeatureMayAlwaysSet& features,
       const Position* MT_NULLABLE position);
 
-  void add_extra_trace(const Frame& propagation_frame);
-
   /**
    * Propagate the taint from the callee to the caller.
    *
diff --git a/source/TaintTree.cpp b/source/TaintTree.cpp
index ee0532f8..b5548dcc 100644
--- a/source/TaintTree.cpp
+++ b/source/TaintTree.cpp
@@ -73,13 +73,6 @@ void TaintTree::add_locally_inferred_features_and_local_position(
   });
 }
 
-void TaintTree::add_extra_trace(const Frame& propagation_frame) {
-  tree_.map([&propagation_frame](Taint taint) {
-    taint.add_extra_trace(propagation_frame);
-    return taint;
-  });
-}
-
 void TaintTree::attach_position(const Position* position) {
   tree_.map(
       [position](Taint taint) { return taint.attach_position(position); });
@@ -122,6 +115,7 @@ void TaintTree::update_maximum_collapse_depth(CollapseDepth collapse_depth) {
 }
 
 void TaintTree::update_with_propagation_trace(const Frame& propagation_frame) {
+  mt_assert(propagation_frame.call_info().is_propagation_with_trace());
   tree_.map([&propagation_frame](const Taint& taint) {
     return taint.update_with_propagation_trace(propagation_frame);
   });
diff --git a/source/TaintTree.h b/source/TaintTree.h
index e203cfa1..c58ef207 100644
--- a/source/TaintTree.h
+++ b/source/TaintTree.h
@@ -107,8 +107,6 @@ class TaintTree final : public sparta::AbstractDomain<TaintTree> {
       const FeatureMayAlwaysSet& features,
       const Position* MT_NULLABLE position);
 
-  void add_extra_trace(const Frame& propagation_frame);
-
   void attach_position(const Position* position);
 
   /* Return all taint in the tree. */
diff --git a/source/tests/integration/end-to-end/code/taint_transforms/expected_output.json b/source/tests/integration/end-to-end/code/taint_transforms/expected_output.json
index 9942c3b2..982151b6 100644
--- a/source/tests/integration/end-to-end/code/taint_transforms/expected_output.json
+++ b/source/tests/integration/end-to-end/code/taint_transforms/expected_output.json
@@ -988,6 +988,20 @@
             {
               "call_info" : "PropagationWithTrace:CallSite",
               "distance" : 2,
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 33,
+                      "path" : "ExtraTraces.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T2@T1:LocalReturn",
               "origins" : 
               [
@@ -1001,6 +1015,20 @@
             {
               "call_info" : "PropagationWithTrace:CallSite",
               "distance" : 2,
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 33,
+                      "path" : "ExtraTraces.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T2@T2:LocalReturn",
               "origins" : 
               [
@@ -1355,6 +1383,20 @@
             {
               "call_info" : "CallSite",
               "distance" : 1,
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T1:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 65,
+                      "path" : "ExtraTraces.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T1@Sink",
               "origins" : 
               [
@@ -1364,6 +1406,20 @@
             {
               "call_info" : "CallSite",
               "distance" : 1,
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 68,
+                      "path" : "ExtraTraces.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T2@Sink",
               "origins" : 
               [
@@ -1880,7 +1936,7 @@
             {
               "call_info" : "CallSite",
               "distance" : 1,
-              "kind" : "T1@T2:Sink",
+              "kind" : "T2:Sink",
               "origins" : 
               [
                 "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V"
@@ -1889,7 +1945,21 @@
             {
               "call_info" : "CallSite",
               "distance" : 1,
-              "kind" : "T2:Sink",
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T1:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 159,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                }
+              ],
+              "kind" : "T1@T2:Sink",
               "origins" : 
               [
                 "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V"
@@ -2315,6 +2385,20 @@
             },
             {
               "call_info" : "Origin",
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T1:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 135,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T1@Sink",
               "origins" : 
               [
@@ -2364,6 +2448,20 @@
           [
             {
               "call_info" : "Origin",
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 142,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                }
+              ],
               "kind" : "T2@Sink",
               "origins" : 
               [
@@ -2413,7 +2511,21 @@
           [
             {
               "call_info" : "Origin",
-              "kind" : "T1:T2@Sink",
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 151,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                }
+              ],
+              "kind" : "T2@Sink",
               "origins" : 
               [
                 "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V"
@@ -2421,7 +2533,32 @@
             },
             {
               "call_info" : "Origin",
-              "kind" : "T2@Sink",
+              "extra_traces" : 
+              [
+                {
+                  "kind" : "T1:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 150,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                },
+                {
+                  "kind" : "T2:LocalReturn",
+                  "origin" : 
+                  {
+                    "position" : 
+                    {
+                      "line" : 151,
+                      "path" : "TaintTransforms.java"
+                    }
+                  }
+                }
+              ],
+              "kind" : "T1:T2@Sink",
               "origins" : 
               [
                 "Lcom/facebook/marianatrench/integrationtests/Origin;.sink:(Ljava/lang/Object;)V"