From ef534d84797d56adeb7a712239afb676e01aeab4 Mon Sep 17 00:00:00 2001
From: Khan Winter <35942988+thecoolwinter@users.noreply.github.com>
Date: Tue, 6 Feb 2024 00:57:58 -0600
Subject: [PATCH] Fix editor restoration and setup

---
 .../Models/CEWorkspaceFileManager.swift       |  4 +-
 .../EditorLayout+StateRestoration.swift       | 48 ++++++++++++-------
 .../Features/Editor/Models/EditorLayout.swift | 17 +++++++
 3 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/CodeEdit/Features/CEWorkspace/Models/CEWorkspaceFileManager.swift b/CodeEdit/Features/CEWorkspace/Models/CEWorkspaceFileManager.swift
index 09ed8554a..56453a092 100644
--- a/CodeEdit/Features/CEWorkspace/Models/CEWorkspaceFileManager.swift
+++ b/CodeEdit/Features/CEWorkspace/Models/CEWorkspaceFileManager.swift
@@ -96,7 +96,9 @@ final class CEWorkspaceFileManager {
         if let file = flattenedFileItems[path] {
             return file
         } else if createIfNotFound {
-            let url = URL(fileURLWithPath: path, relativeTo: folderUrl)
+            guard let url = URL(string: path, relativeTo: folderUrl) else {
+                return nil
+            }
 
             // Drill down towards the file, indexing any directories needed. If file is not in the `folderURL` or
             // subdirectories, exit.
diff --git a/CodeEdit/Features/Editor/Models/EditorLayout+StateRestoration.swift b/CodeEdit/Features/Editor/Models/EditorLayout+StateRestoration.swift
index 0b48fb900..8279986f1 100644
--- a/CodeEdit/Features/Editor/Models/EditorLayout+StateRestoration.swift
+++ b/CodeEdit/Features/Editor/Models/EditorLayout+StateRestoration.swift
@@ -13,22 +13,38 @@ extension EditorManager {
     /// Restores the tab manager from a captured state obtained using `saveRestorationState`
     /// - Parameter workspace: The workspace to retrieve state from.
     func restoreFromState(_ workspace: WorkspaceDocument) {
-        guard let fileManager = workspace.workspaceFileManager,
-              let data = workspace.getFromWorkspaceState(.openTabs) as? Data,
-              let state = try? JSONDecoder().decode(EditorRestorationState.self, from: data) else {
-            return
-        }
+        do {
+            guard let fileManager = workspace.workspaceFileManager,
+                  let data = workspace.getFromWorkspaceState(.openTabs) as? Data else {
+                return
+            }
 
-        guard !state.groups.isEmpty else {
-            logger.warning("Empty Editor State found, restoring to clean editor state.")
-            initCleanState()
-            return
-        }
+            let state = try JSONDecoder().decode(EditorRestorationState.self, from: data)
+
+            guard !state.groups.isEmpty else {
+                logger.warning("Empty Editor State found, restoring to clean editor state.")
+                initCleanState()
+                return
+            }
 
-        fixRestoredEditorLayout(state.groups, fileManager: fileManager)
-        self.editorLayout = state.groups
-        self.activeEditor = activeEditor
-        switchToActiveEditor()
+            guard let activeEditor = state.groups.find(
+                editor: state.activeEditor
+            ) ?? state.groups.findSomeEditor() else {
+                logger.warning("Editor state could not restore active editor.")
+                initCleanState()
+                return
+            }
+
+            fixRestoredEditorLayout(state.groups, fileManager: fileManager)
+
+            self.editorLayout = state.groups
+            self.activeEditor = activeEditor
+            switchToActiveEditor()
+        } catch {
+            logger.warning(
+                "Could not restore editor state from saved data: \(error.localizedDescription, privacy: .public)"
+            )
+        }
     }
 
     /// Fix any hanging files after restoring from saved state.
@@ -89,7 +105,7 @@ extension EditorManager {
 
     func saveRestorationState(_ workspace: WorkspaceDocument) {
         if let data = try? JSONEncoder().encode(
-            EditorRestorationState(focus: activeEditor, groups: editorLayout)
+            EditorRestorationState(activeEditor: activeEditor.id, groups: editorLayout)
         ) {
             workspace.addToWorkspaceState(key: .openTabs, value: data)
         } else {
@@ -99,7 +115,7 @@ extension EditorManager {
 }
 
 struct EditorRestorationState: Codable {
-    var focus: Editor
+    var activeEditor: UUID
     var groups: EditorLayout
 }
 
diff --git a/CodeEdit/Features/Editor/Models/EditorLayout.swift b/CodeEdit/Features/Editor/Models/EditorLayout.swift
index b43e663f9..86932561a 100644
--- a/CodeEdit/Features/Editor/Models/EditorLayout.swift
+++ b/CodeEdit/Features/Editor/Models/EditorLayout.swift
@@ -44,6 +44,23 @@ enum EditorLayout: Equatable {
         }
     }
 
+    func find(editor id: UUID) -> Editor? {
+        switch self {
+        case .one(let editor):
+            if editor.id == id {
+                return editor
+            }
+        case .vertical(let splitViewData), .horizontal(let splitViewData):
+            for layout in splitViewData.editorLayouts {
+                if let editor = layout.find(editor: id) {
+                    return editor
+                }
+            }
+        }
+
+        return nil
+    }
+
     /// Forms a set of all files currently represented by tabs.
     func gatherOpenFiles() -> Set<CEWorkspaceFile> {
         switch self {