diff --git a/libs/input/SpriteController.cpp b/libs/input/SpriteController.cpp
index 0c6dac0..da21438 100644
--- a/libs/input/SpriteController.cpp
+++ b/libs/input/SpriteController.cpp
@@ -27,10 +27,12 @@
 
 // --- SpriteController ---
 
-SpriteController::SpriteController(const sp<Looper>& looper, int32_t overlayLayer) :
-        mLooper(looper), mOverlayLayer(overlayLayer) {
+SpriteController::SpriteController(const sp<Looper>& looper, int32_t overlayLayer,
+                                   ParentSurfaceProvider parentSurfaceProvider)
+      : mLooper(looper),
+        mOverlayLayer(overlayLayer),
+        mParentSurfaceProvider(std::move(parentSurfaceProvider)) {
     mHandler = new WeakMessageHandler(this);
-
     mLocked.transactionNestingCount = 0;
     mLocked.deferredSpriteUpdate = false;
 }
@@ -168,8 +170,7 @@
 
         // If surface is a new one, we have to set right layer stack.
         if (update.surfaceChanged || update.state.dirty & DIRTY_DISPLAY_ID) {
-            t.setLayerStack(update.state.surfaceControl,
-                            ui::LayerStack::fromValue(update.state.displayId));
+            t.reparent(update.state.surfaceControl, mParentSurfaceProvider(update.state.displayId));
             needApplyTransaction = true;
         }
     }
diff --git a/libs/input/SpriteController.h b/libs/input/SpriteController.h
index 137b564..2a80d95 100644
--- a/libs/input/SpriteController.h
+++ b/libs/input/SpriteController.h
@@ -114,7 +114,8 @@
     virtual ~SpriteController();
 
 public:
-    SpriteController(const sp<Looper>& looper, int32_t overlayLayer);
+    using ParentSurfaceProvider = std::function<sp<SurfaceControl>(int /*displayId*/)>;
+    SpriteController(const sp<Looper>& looper, int32_t overlayLayer, ParentSurfaceProvider parent);
 
     /* Creates a new sprite, initially invisible. */
     virtual sp<Sprite> createSprite();
@@ -245,6 +246,7 @@
     sp<Looper> mLooper;
     const int32_t mOverlayLayer;
     sp<WeakMessageHandler> mHandler;
+    ParentSurfaceProvider mParentSurfaceProvider;
 
     sp<SurfaceComposerClient> mSurfaceComposerClient;
 
diff --git a/libs/input/tests/mocks/MockSpriteController.h b/libs/input/tests/mocks/MockSpriteController.h
index a034f66..62f1d65 100644
--- a/libs/input/tests/mocks/MockSpriteController.h
+++ b/libs/input/tests/mocks/MockSpriteController.h
@@ -26,7 +26,8 @@
 class MockSpriteController : public SpriteController {
 
 public:
-    MockSpriteController(sp<Looper> looper) : SpriteController(looper, 0) {}
+    MockSpriteController(sp<Looper> looper)
+          : SpriteController(looper, 0, [](int) { return nullptr; }) {}
     ~MockSpriteController() {}
 
     MOCK_METHOD(sp<Sprite>, createSprite, (), (override));