diff --git a/services/surfaceflinger/LayerProtoHelper.cpp b/services/surfaceflinger/LayerProtoHelper.cpp
index 0d2987c..885c3d3 100644
--- a/services/surfaceflinger/LayerProtoHelper.cpp
+++ b/services/surfaceflinger/LayerProtoHelper.cpp
@@ -263,7 +263,7 @@
     outRegion.bottom = proto.bottom();
 }
 
-perfetto::protos::LayersProto LayerProtoFromSnapshotGenerator::generate(
+LayerProtoFromSnapshotGenerator& LayerProtoFromSnapshotGenerator::with(
         const frontend::LayerHierarchy& root) {
     mLayersProto.clear_layers();
     mVisitedLayers.clear();
@@ -305,9 +305,40 @@
         }
     }
 
-    mDefaultSnapshots.clear();
-    mChildToRelativeParent.clear();
-    return std::move(mLayersProto);
+    return *this;
+}
+
+LayerProtoFromSnapshotGenerator& LayerProtoFromSnapshotGenerator::withOffscreenLayers(
+        const frontend::LayerHierarchy& offscreenRoot) {
+    // Add a fake invisible root layer to the proto output and parent all the offscreen layers to
+    // it.
+    perfetto::protos::LayerProto* rootProto = mLayersProto.add_layers();
+    const int32_t offscreenRootLayerId = INT32_MAX - 2;
+    rootProto->set_id(offscreenRootLayerId);
+    rootProto->set_name("Offscreen Root");
+    rootProto->set_parent(-1);
+
+    perfetto::protos::LayersProto offscreenLayers =
+            LayerProtoFromSnapshotGenerator(mSnapshotBuilder, mDisplayInfos, mLegacyLayers,
+                                            mTraceFlags)
+                    .with(offscreenRoot)
+                    .generate();
+
+    for (int i = 0; i < offscreenLayers.layers_size(); i++) {
+        perfetto::protos::LayerProto* layerProto = offscreenLayers.mutable_layers()->Mutable(i);
+        if (layerProto->parent() == -1) {
+            layerProto->set_parent(offscreenRootLayerId);
+            // Add layer as child of the fake root
+            rootProto->add_children(layerProto->id());
+        }
+    }
+
+    mLayersProto.mutable_layers()->Reserve(mLayersProto.layers_size() +
+                                           offscreenLayers.layers_size());
+    std::copy(offscreenLayers.layers().begin(), offscreenLayers.layers().end(),
+              RepeatedFieldBackInserter(mLayersProto.mutable_layers()));
+
+    return *this;
 }
 
 frontend::LayerSnapshot* LayerProtoFromSnapshotGenerator::getSnapshot(
diff --git a/services/surfaceflinger/LayerProtoHelper.h b/services/surfaceflinger/LayerProtoHelper.h
index d672012..c0198b6 100644
--- a/services/surfaceflinger/LayerProtoHelper.h
+++ b/services/surfaceflinger/LayerProtoHelper.h
@@ -88,7 +88,12 @@
             mLegacyLayers(legacyLayers),
             mDisplayInfos(displayInfos),
             mTraceFlags(traceFlags) {}
-    perfetto::protos::LayersProto generate(const frontend::LayerHierarchy& root);
+    LayerProtoFromSnapshotGenerator& with(const frontend::LayerHierarchy& root);
+    // Creates a fake root and adds all offscreen layers from the passed in hierarchy to the fake
+    // root
+    LayerProtoFromSnapshotGenerator& withOffscreenLayers(
+            const frontend::LayerHierarchy& offscreenRoot);
+    perfetto::protos::LayersProto generate() { return mLayersProto; };
 
 private:
     void writeHierarchyToProto(const frontend::LayerHierarchy& root,
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 7d2384e..ab10ee4 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -5867,9 +5867,16 @@
         }
     }
 
-    return LayerProtoFromSnapshotGenerator(mLayerSnapshotBuilder, mFrontEndDisplayInfos,
-                                           mLegacyLayers, traceFlags)
-            .generate(mLayerHierarchyBuilder.getHierarchy());
+    auto traceGenerator =
+            LayerProtoFromSnapshotGenerator(mLayerSnapshotBuilder, mFrontEndDisplayInfos,
+                                            mLegacyLayers, traceFlags)
+                    .with(mLayerHierarchyBuilder.getHierarchy());
+
+    if (traceFlags & LayerTracing::Flag::TRACE_EXTRA) {
+        traceGenerator.withOffscreenLayers(mLayerHierarchyBuilder.getOffscreenHierarchy());
+    }
+
+    return traceGenerator.generate();
 }
 
 google::protobuf::RepeatedPtrField<perfetto::protos::DisplayProto>
@@ -5903,36 +5910,6 @@
     getHwComposer().dump(result);
 }
 
-void SurfaceFlinger::dumpOffscreenLayersProto(perfetto::protos::LayersProto& layersProto,
-                                              uint32_t traceFlags) const {
-    // Add a fake invisible root layer to the proto output and parent all the offscreen layers to
-    // it.
-    perfetto::protos::LayerProto* rootProto = layersProto.add_layers();
-    const int32_t offscreenRootLayerId = INT32_MAX - 2;
-    rootProto->set_id(offscreenRootLayerId);
-    rootProto->set_name("Offscreen Root");
-    rootProto->set_parent(-1);
-
-    perfetto::protos::LayersProto offscreenLayers =
-            LayerProtoFromSnapshotGenerator(mLayerSnapshotBuilder, mFrontEndDisplayInfos,
-                                            mLegacyLayers, traceFlags)
-                    .generate(mLayerHierarchyBuilder.getOffscreenHierarchy());
-
-    for (int i = 0; i < offscreenLayers.layers_size(); i++) {
-        perfetto::protos::LayerProto* layerProto = offscreenLayers.mutable_layers()->Mutable(i);
-        if (layerProto->parent() == -1) {
-            layerProto->set_parent(offscreenRootLayerId);
-            // Add layer as child of the fake root
-            rootProto->add_children(layerProto->id());
-        }
-    }
-
-    layersProto.mutable_layers()->Reserve(layersProto.layers_size() +
-                                          offscreenLayers.layers_size());
-    std::copy(offscreenLayers.layers().begin(), offscreenLayers.layers().end(),
-              RepeatedFieldBackInserter(layersProto.mutable_layers()));
-}
-
 perfetto::protos::LayersProto SurfaceFlinger::dumpProtoFromMainThread(uint32_t traceFlags) {
     return mScheduler
             ->schedule([=, this]() FTL_FAKE_GUARD(kMainThreadContext) {
@@ -8436,9 +8413,6 @@
                                             0);
 
     auto layers = dumpDrawingStateProto(traceFlags);
-    if (traceFlags & LayerTracing::Flag::TRACE_EXTRA) {
-        dumpOffscreenLayersProto(layers);
-    }
     *snapshot.mutable_layers() = std::move(layers);
 
     if (traceFlags & LayerTracing::Flag::TRACE_HWC) {
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 282c8cf..873fac2 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -1128,9 +1128,6 @@
 
     perfetto::protos::LayersProto dumpDrawingStateProto(uint32_t traceFlags) const
             REQUIRES(kMainThreadContext);
-    void dumpOffscreenLayersProto(perfetto::protos::LayersProto& layersProto,
-                                  uint32_t traceFlags = LayerTracing::TRACE_ALL) const
-            REQUIRES(kMainThreadContext);
     google::protobuf::RepeatedPtrField<perfetto::protos::DisplayProto> dumpDisplayProto() const;
     void doActiveLayersTracingIfNeeded(bool isCompositionComputed, bool visibleRegionDirty,
                                        TimePoint, VsyncId) REQUIRES(kMainThreadContext);
diff --git a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
index 617ea2c..1dba175 100644
--- a/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
+++ b/services/surfaceflinger/Tracing/tools/LayerTraceGenerator.cpp
@@ -162,7 +162,10 @@
 
         auto layersProto =
                 LayerProtoFromSnapshotGenerator(snapshotBuilder, displayInfos, {}, traceFlags)
-                        .generate(hierarchyBuilder.getHierarchy());
+                        .with(hierarchyBuilder.getHierarchy())
+                        .withOffscreenLayers(hierarchyBuilder.getOffscreenHierarchy())
+                        .generate();
+
         auto displayProtos = LayerProtoHelper::writeDisplayInfoToProto(displayInfos);
         if (!onlyLastEntry || (i == traceFile.entry_size() - 1)) {
             perfetto::protos::LayersSnapshotProto snapshotProto{};