@@ -120,14 +120,10 @@ void TopLevelWindowModel::setSurfaceManager(unityapi::SurfaceManagerInterface *s
120120 m_surfaceManager = surfaceManager;
121121
122122 if (m_surfaceManager) {
123+ connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::surfacesAddedToWorkspace, this, &TopLevelWindowModel::onSurfacesAddedToWorkspace);
123124 connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::surfacesRaised, this, &TopLevelWindowModel::onSurfacesRaised);
124125 connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::modificationsStarted, this, &TopLevelWindowModel::onModificationsStarted);
125126 connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::modificationsEnded, this, &TopLevelWindowModel::onModificationsEnded);
126−
127− connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::surfacesAddedToWorkspace,
128− this, &TopLevelWindowModel::onSurfacesAddedToWorkspace);
129− connect(m_surfaceManager, &unityapi::SurfaceManagerInterface::surfacesAboutToBeRemovedFromWorkspace,
130− this, &TopLevelWindowModel::onSurfacesAboutToBeRemovedFromWorkspace);
131127 }
132128
133129 refreshWindows();
@@ -382,7 +378,10 @@ void TopLevelWindowModel::onSurfacesAddedToWorkspace(const std::shared_ptr<miral
382378 const QVector<unity::shell::application::MirSurfaceInterface*> surfaces)
383379{
384380 if (!m_workspace || !m_applicationManager) return;
385− if (workspace != m_workspace->workspace()) return;
381+ if (workspace != m_workspace->workspace()) {
382+ removeSurfaces(surfaces);
383+ return;
384+ }
386385
387386 Q_FOREACH(auto surface, surfaces) {
388387 if (m_allSurfaces.contains(surface)) continue;
@@ -426,11 +425,8 @@ void TopLevelWindowModel::onSurfacesAddedToWorkspace(const std::shared_ptr<miral
426425 }
427426}
428427
429−void TopLevelWindowModel::onSurfacesAboutToBeRemovedFromWorkspace(const std::shared_ptr<miral::Workspace> &workspace,
430− const QVector<unity::shell::application::MirSurfaceInterface *> surfaces)
428+void TopLevelWindowModel::removeSurfaces(const QVector<unity::shell::application::MirSurfaceInterface *> surfaces)
431429{
432− if (!m_workspace || workspace != m_workspace->workspace()) return;
433−
434430 int start = -1;
435431 int end = -1;
436432 for (auto iter = surfaces.constBegin(); iter != surfaces.constEnd();) {
@@ -440,7 +436,7 @@ void TopLevelWindowModel::onSurfacesAboutToBeRemovedFromWorkspace(const std::sha
440436 // Do removals in adjacent blocks.
441437 start = end = indexOf(surface);
442438 if (start == -1) {
443− // probably a child surface
439+ // could be a child surface
444440 m_allSurfaces.remove(surface);
445441 continue;
446442 }
@@ -538,6 +534,15 @@ void TopLevelWindowModel::setInputMethodWindow(Window *window)
538534void TopLevelWindowModel::removeInputMethodWindow()
539535{
540536 if (m_inputMethodWindow) {
537+ auto surface = m_inputMethodWindow->surface();
538+ if (surface) {
539+ m_allSurfaces.remove(surface);
540+ }
541+ if (m_focusedWindow == m_inputMethodWindow) {
542+ setFocusedWindow(nullptr);
543+ m_focusedWindowCleared = false;
544+ }
545+
541546 delete m_inputMethodWindow;
542547 m_inputMethodWindow = nullptr;
543548 Q_EMIT inputMethodSurfaceChanged(nullptr);
@@ -549,7 +554,7 @@ void TopLevelWindowModel::onSurfacesRaised(const QVector<unityapi::MirSurfaceInt
549554 DEBUG_MSG << "(" << surfaces << ")";
550555 const int raiseCount = surfaces.size();
551556 for (int i = 0; i < raiseCount; i++) {
552− int fromIndex = findIndexOf(surfaces[i]);
557+ int fromIndex = indexOf(surfaces[i]);
553558 if (fromIndex != -1) {
554559 move(fromIndex, 0);
555560 }
@@ -576,16 +581,6 @@ QVariant TopLevelWindowModel::data(const QModelIndex& index, int role) const
576581 }
577582}
578583
579−int TopLevelWindowModel::findIndexOf(const unityapi::MirSurfaceInterface *surface) const
580−{
581− for (int i=0; i<m_windowModel.count(); i++) {
582− if (m_windowModel[i].window->surface() == surface) {
583− return i;
584− }
585− }
586− return -1;
587−}
588−
589584int TopLevelWindowModel::generateId()
590585{
591586 int id = m_nextId;